diff --git a/pkg/services/datasources/fakes/fake_datasource_service.go b/pkg/services/datasources/fakes/fake_datasource_service.go index 19e7bbb43a5..7637657f7cf 100644 --- a/pkg/services/datasources/fakes/fake_datasource_service.go +++ b/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 } diff --git a/pkg/services/ngalert/writer/datasourcewriter.go b/pkg/services/ngalert/writer/datasourcewriter.go index 19af0db0a5b..cc7036974e7 100644 --- a/pkg/services/ngalert/writer/datasourcewriter.go +++ b/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 diff --git a/pkg/services/ngalert/writer/datasourcewriter_test.go b/pkg/services/ngalert/writer/datasourcewriter_test.go index f06d2d95de4..349d8b9191a 100644 --- a/pkg/services/ngalert/writer/datasourcewriter_test.go +++ b/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() diff --git a/pkg/services/ngalert/writer/testing.go b/pkg/services/ngalert/writer/testing.go index 91b1bdb6b10..8b764657d6e 100644 --- a/pkg/services/ngalert/writer/testing.go +++ b/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()