InfluxDB: Add custom metadata to the first frame via header (#88698)

* add custom metadata via header to the response

* use strings.CutPrefix instead of regex

pkg/tsdb/influxdb/influxql/influxql.go

@@ -23,7 +23,10 @@ import (
 	"github.com/grafana/grafana/pkg/tsdb/influxdb/models"
 )
 
-const defaultRetentionPolicy = "default"
+const (
+	defaultRetentionPolicy = "default"
+	metadataPrefix         = "x-grafana-meta-add-"
+)
 
 var (
 	ErrInvalidHttpMode = errors.New("'httpMode' should be either 'GET' or 'POST'")
@@ -195,9 +198,27 @@ func execute(ctx context.Context, tracer trace.Tracer, dsInfo *models.Datasource
 	} else {
 		resp = buffered.ResponseParse(res.Body, res.StatusCode, query)
 	}
+
+	if resp.Frames != nil && len(resp.Frames) > 0 {
+		resp.Frames[0].Meta.Custom = readCustomMetadata(res)
+	}
+
 	return *resp, nil
 }
 
+func readCustomMetadata(res *http.Response) map[string]any {
+	var result map[string]any
+	for k := range res.Header {
+		if key, found := strings.CutPrefix(strings.ToLower(k), metadataPrefix); found {
+			if result == nil {
+				result = make(map[string]any)
+			}
+			result[key] = res.Header.Get(k)
+		}
+	}
+	return result
+}
+
 // startTrace setups a trace but does not panic if tracer is nil which helps with testing
 func startTrace(ctx context.Context, tracer trace.Tracer, name string, attributes ...attribute.KeyValue) (context.Context, func()) {
 	if tracer == nil {

pkg/tsdb/influxdb/influxql/influxql_test.go

@@ -3,6 +3,7 @@ package influxql
 import (
 	"context"
 	"io"
+	"net/http"
 	"net/url"
 	"testing"
 
@@ -60,3 +61,52 @@ func TestExecutor_createRequest(t *testing.T) {
 		require.EqualError(t, err, ErrInvalidHttpMode.Error())
 	})
 }
+
+func TestReadCustomMetadata(t *testing.T) {
+	t.Run("should read nothing if no X-Grafana-Meta-Add-<Thing> header exists", func(t *testing.T) {
+		header := http.Header{}
+		header.Add("content-type", "text/html")
+		header.Add("content-encoding", "gzip")
+		res := &http.Response{
+			Header: header,
+		}
+		result := readCustomMetadata(res)
+		require.Nil(t, result)
+	})
+
+	t.Run("should read X-Grafana-Meta-Add-<Thing> header", func(t *testing.T) {
+		header := http.Header{}
+		header.Add("content-type", "text/html")
+		header.Add("content-encoding", "gzip")
+		header.Add("X-Grafana-Meta-Add-TestThing", "test1234")
+		res := &http.Response{
+			Header: header,
+		}
+		result := readCustomMetadata(res)
+		expected := map[string]any{
+			"testthing": "test1234",
+		}
+		require.NotNil(t, result)
+		require.Equal(t, expected, result)
+	})
+
+	t.Run("should read multiple X-Grafana-Meta-Add-<Thing> header", func(t *testing.T) {
+		header := http.Header{}
+		header.Add("content-type", "text/html")
+		header.Add("content-encoding", "gzip")
+		header.Add("X-Grafana-Meta-Add-TestThing", "test111")
+		header.Add("X-Grafana-Meta-Add-TestThing2", "test222")
+		header.Add("X-Grafana-Meta-Add-Test-Other", "other")
+		res := &http.Response{
+			Header: header,
+		}
+		result := readCustomMetadata(res)
+		expected := map[string]any{
+			"testthing":  "test111",
+			"testthing2": "test222",
+			"test-other": "other",
+		}
+		require.NotNil(t, result)
+		require.Equal(t, expected, result)
+	})
+}