Alerting: Use data source headers when remote writing (#114528)

pkg/services/datasources/fakes/fake_datasource_service.go

@@ -15,6 +15,9 @@ type FakeDataSourceService struct {
 	lastID                int64
 	DataSources           []*datasources.DataSource
 	SimulatePluginFailure bool
+
+	// UID -> Headers
+	DataSourceHeaders map[string]http.Header
 }
 
 var _ datasources.DataSourceService = &FakeDataSourceService{}
@@ -152,5 +155,5 @@ func (s *FakeDataSourceService) DecryptedPassword(ctx context.Context, ds *datas
 }
 
 func (s *FakeDataSourceService) CustomHeaders(ctx context.Context, ds *datasources.DataSource) (http.Header, error) {
-	return nil, nil
+	return s.DataSourceHeaders[ds.UID], nil
 }

pkg/services/ngalert/writer/datasourcewriter.go

@@ -205,11 +205,23 @@ func (w *DatasourceWriter) makeWriter(ctx context.Context, orgID int64, dsUID st
 		return nil, err
 	}
 
+	// We need to add the writer headers (valid for any data source) and any data-source-specific headers.
 	headers := make(http.Header)
 	for k, v := range w.cfg.CustomHeaders {
 		headers.Add(k, v)
 	}
 
+	dsHeaders, err := w.datasources.CustomHeaders(ctx, ds)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get headers for data source: %w", err)
+	}
+
+	for k, values := range dsHeaders {
+		for _, v := range values {
+			headers.Add(k, v)
+		}
+	}
+
 	var backend backendType
 	if dsUID == string(grafanaCloudPromType) {
 		backend = grafanaCloudPromType

pkg/services/ngalert/writer/datasourcewriter_test.go

@@ -56,13 +56,14 @@ func (m *mockHTTPClientProvider) New(options ...sdkhttpclient.Options) (*http.Cl
 type testDataSources struct {
 	dsfakes.FakeDataSourceService
 
-	prom1, prom2, prom3 *TestRemoteWriteTarget
+	prom1, prom2, prom3, prom4 *TestRemoteWriteTarget
 }
 
 func (t *testDataSources) Reset() {
 	t.prom1.Reset()
 	t.prom2.Reset()
 	t.prom3.Reset()
+	t.prom4.Reset()
 }
 
 func setupDataSources(t *testing.T) *testDataSources {
@@ -70,7 +71,9 @@ func setupDataSources(t *testing.T) *testDataSources {
 		prom1: NewTestRemoteWriteTarget(t),
 		prom2: NewTestRemoteWriteTarget(t),
 		prom3: NewTestRemoteWriteTarget(t),
+		prom4: NewTestRemoteWriteTarget(t),
 	}
+	res.DataSourceHeaders = make(map[string]http.Header)
 
 	t.Cleanup(func() {
 		res.prom1.Close()
@@ -81,6 +84,9 @@ func setupDataSources(t *testing.T) *testDataSources {
 	t.Cleanup(func() {
 		res.prom3.Close()
 	})
+	t.Cleanup(func() {
+		res.prom4.Close()
+	})
 
 	p1, _ := res.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
 		Name:     "prom-1",
@@ -107,7 +113,7 @@ func setupDataSources(t *testing.T) *testDataSources {
 		Type: datasources.DS_LOKI,
 	})
 
-	// Add a third Prometheus datasource that uses PDC
+	// Add a third Prometheus datasource that uses PDC.
 	p3, _ := res.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
 		Name: "prom-3",
 		UID:  "prom-3",
@@ -123,6 +129,21 @@ func setupDataSources(t *testing.T) *testDataSources {
 
 	require.True(t, p3.IsSecureSocksDSProxyEnabled())
 
+	// Add a fourth Prometheus datasource with headers in the JSON config.
+	p4, _ := res.AddDataSource(context.Background(), &datasources.AddDataSourceCommand{
+		Name:     "prom-4",
+		UID:      "prom-4",
+		Type:     datasources.DS_PROMETHEUS,
+		JsonData: simplejson.MustJson([]byte(`{"prometheusType":"Prometheus"}`)),
+	})
+	p4.URL = res.prom4.srv.URL
+	res.prom4.ExpectedPath = "/api/v1/write"
+	res.DataSourceHeaders["prom-4"] = http.Header{
+		"X-Scope-OrgID":   []string{"test-user"},
+		"X-Test-Header":   []string{"test-value"},
+		"X-Double-Header": []string{"one", "two", "three"},
+	}
+
 	return res
 }
 
@@ -204,6 +225,45 @@ func TestDatasourceWriter(t *testing.T) {
 		assert.Equal(t, headers[header2], testDS.prom1.LastHeaders.Get(header2))
 	})
 
+	t.Run("when data source headers are configured, they are passed to the request", func(t *testing.T) {
+		testDS.Reset()
+		overwrittenHeader := "X-Test-Header"
+		cHeaders := map[string]string{
+			"X-Custom-Header":  "test-value",
+			"X-Another-Header": "another-value",
+			overwrittenHeader:  "overwritten", // Data source headers should be overwritten by custom headers.
+		}
+
+		cfg = DatasourceWriterConfig{
+			Timeout:              time.Second * 5,
+			DefaultDatasourceUID: "prom-1",
+			CustomHeaders:        cHeaders,
+		}
+		writer = NewDatasourceWriter(cfg, testDS, httpclient.NewProvider(), pluginContextProvider, clock.New(), log.New("test"), met)
+
+		uid := "prom-4"
+		err := writer.WriteDatasource(context.Background(), uid, "metric", time.Now(), frames, 1, map[string]string{})
+		require.NoError(t, err)
+
+		dsHeaders := testDS.DataSourceHeaders[uid]
+		require.Len(t, dsHeaders, 3)
+
+		// We're confirming we have a data source header with the same name but different value.
+		// This one should not be sent in the request.
+		require.NotEmpty(t, dsHeaders[overwrittenHeader])
+		require.NotEqual(t, dsHeaders[overwrittenHeader], cHeaders[overwrittenHeader])
+
+		// All headers (except for the one that was overwritten) should have been used.
+		for k, vv := range dsHeaders {
+			if k != overwrittenHeader {
+				assert.Equal(t, vv, testDS.prom4.LastHeaders.Values(k))
+			}
+		}
+		for k, v := range cHeaders {
+			assert.Equal(t, v, testDS.prom4.LastHeaders.Get(k))
+		}
+	})
+
 	t.Run("when PDC is enabled proxy options are passed to HTTP client provider", func(t *testing.T) {
 		testDS.Reset()
 

pkg/services/ngalert/writer/testing.go

@@ -1,6 +1,7 @@
 package writer
 
 import (
+	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -37,7 +38,7 @@ func NewTestRemoteWriteTarget(t *testing.T) *TestRemoteWriteTarget {
 
 	handler := func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path != target.ExpectedPath {
-			require.Fail(t, "Received unexpected request for endpoint %s", r.URL.Path)
+			require.Fail(t, fmt.Sprintf("Received unexpected request for endpoint %s", r.URL.Path))
 		}
 
 		target.mtx.Lock()