feat: grafana query metadata (#1743)

Co-authored-by: Andrew Hackmann <5140848+bossinc@users.noreply.github.com>

Author: MattiasMTS

pkg/plugin/driver.go

@@ -31,6 +31,16 @@ import (
 	"golang.org/x/net/proxy"
 )
 
+type grafanaHeadersKeyType struct{}
+
+var grafanaHeadersKey = grafanaHeadersKeyType{}
+
+type grafanaHeaders struct {
+	DashboardUID string
+	PanelID      string
+	RuleUID      string
+}
+
 // Clickhouse defines how to connect to a Clickhouse datasource
 type Clickhouse struct {
 	SchemaDatasource *schemas.SchemaDatasource
@@ -125,14 +135,19 @@ func CheckMinServerVersion(conn *sql.DB, major, minor, patch uint64) (bool, erro
 			version.Patch, _ = strconv.ParseUint(v, 10, 64)
 		}
 	}
-	if version.Major < major || (version.Major == major && version.Minor < minor) || (version.Major == major && version.Minor == minor && version.Patch < patch) {
+	if version.Major < major || (version.Major == major && version.Minor < minor) ||
+		(version.Major == major && version.Minor == minor && version.Patch < patch) {
 		return false, nil
 	}
 	return true, nil
 }
 
 // Connect opens a sql.DB connection using datasource settings
-func (h *Clickhouse) Connect(ctx context.Context, config backend.DataSourceInstanceSettings, message json.RawMessage) (*sql.DB, error) {
+func (h *Clickhouse) Connect(
+	ctx context.Context,
+	config backend.DataSourceInstanceSettings,
+	message json.RawMessage,
+) (*sql.DB, error) {
 	ctx, span := tracing.DefaultTracer().Start(ctx, "clickhouse connect", trace.WithAttributes(
 		attribute.String("db.system", "clickhouse"),
 	))
@@ -354,7 +369,22 @@ func (h *Clickhouse) Settings(ctx context.Context, config backend.DataSourceInst
 	}
 }
 
-func (h *Clickhouse) MutateQueryData(ctx context.Context, req *backend.QueryDataRequest) (context.Context, *backend.QueryDataRequest) {
+// MutateQueryData extracts Grafana contextual headers from the request and
+// stores them in the context for ClickHouse query metadata injection.
+func (h *Clickhouse) MutateQueryData(
+	ctx context.Context,
+	req *backend.QueryDataRequest,
+) (context.Context, *backend.QueryDataRequest) {
+	headers := req.GetHTTPHeaders()
+	gh := grafanaHeaders{
+		DashboardUID: headers.Get("X-Dashboard-Uid"),
+		PanelID:      headers.Get("X-Panel-Id"),
+		RuleUID:      headers.Get("X-Rule-Uid"),
+	}
+	if gh.DashboardUID != "" || gh.PanelID != "" || gh.RuleUID != "" {
+		ctx = context.WithValue(ctx, grafanaHeadersKey, gh)
+	}
+
 	req = preprocessGrafanaSQL(req)
 	return ctx, req
 }
@@ -415,10 +445,28 @@ func (h *Clickhouse) MutateQuery(ctx context.Context, req backend.DataQuery) (co
 
 	defer span.End()
 
+	comments := make([]string, 0, 4)
+
 	if user := backend.UserFromContext(ctx); user != nil {
+		comments = append(comments, "grafana_user:"+user.Login)
+	}
+
+	if gh, ok := ctx.Value(grafanaHeadersKey).(grafanaHeaders); ok {
+		if gh.DashboardUID != "" {
+			comments = append(comments, "grafana_dashboard:"+gh.DashboardUID)
+		}
+		if gh.PanelID != "" {
+			comments = append(comments, "grafana_panel:"+gh.PanelID)
+		}
+		if gh.RuleUID != "" {
+			comments = append(comments, "grafana_rule:"+gh.RuleUID)
+		}
+	}
+
+	if len(comments) > 0 {
 		ctx = clickhouse.Context(ctx, clickhouse.WithClientInfo(clickhouse.ClientInfo{
 			Products: nil,
-			Comment:  []string{fmt.Sprintf("grafana_user:%s", user.Login)},
+			Comment:  comments,
 		}))
 	}
 

pkg/plugin/driver_test.go

@@ -1,12 +1,14 @@
 package plugin
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"testing"
 
 	"github.com/ClickHouse/clickhouse-go/v2"
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
 	"github.com/stretchr/testify/assert"
 )
@@ -317,3 +319,109 @@ func TestContainsClickHouseException(t *testing.T) {
 		assert.True(t, result)
 	})
 }
+
+func TestMutateQueryData(t *testing.T) {
+	h := &Clickhouse{}
+
+	tests := []struct {
+		name   string
+		headers map[string]string
+		want   grafanaHeaders
+		stored bool
+	}{
+		{
+			name: "all headers",
+			headers: map[string]string{
+				"http_X-Dashboard-Uid": "dash-abc123",
+				"http_X-Panel-Id":      "42",
+				"http_X-Rule-Uid":      "rule-xyz",
+			},
+			want:   grafanaHeaders{DashboardUID: "dash-abc123", PanelID: "42", RuleUID: "rule-xyz"},
+			stored: true,
+		},
+		{
+			name:    "empty headers",
+			headers: map[string]string{},
+			stored:  false,
+		},
+		{
+			name:    "only dashboard",
+			headers: map[string]string{"http_X-Dashboard-Uid": "dash-only"},
+			want:    grafanaHeaders{DashboardUID: "dash-only"},
+			stored:  true,
+		},
+		{
+			name:    "only panel",
+			headers: map[string]string{"http_X-Panel-Id": "99"},
+			want:    grafanaHeaders{PanelID: "99"},
+			stored:  true,
+		},
+		{
+			name:    "only rule",
+			headers: map[string]string{"http_X-Rule-Uid": "alert-rule-1"},
+			want:    grafanaHeaders{RuleUID: "alert-rule-1"},
+			stored:  true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := &backend.QueryDataRequest{Headers: tt.headers}
+			newCtx, _ := h.MutateQueryData(t.Context(), req)
+
+			gh, ok := newCtx.Value(grafanaHeadersKey).(grafanaHeaders)
+			assert.Equal(t, tt.stored, ok)
+			if tt.stored {
+				assert.Equal(t, tt.want, gh)
+			}
+		})
+	}
+
+	t.Run("nil headers does not panic", func(t *testing.T) {
+		newCtx, newReq := h.MutateQueryData(t.Context(), &backend.QueryDataRequest{})
+		assert.NotNil(t, newCtx)
+		assert.NotNil(t, newReq)
+	})
+}
+
+func TestMutateQuery_GrafanaMetadata(t *testing.T) {
+	h := &Clickhouse{}
+
+	t.Run("includes dashboard and panel from context", func(t *testing.T) {
+		ctx := context.WithValue(t.Context(), grafanaHeadersKey, grafanaHeaders{
+			DashboardUID: "my-dashboard",
+			PanelID:      "7",
+			RuleUID:      "alert-1",
+		})
+
+		newCtx, _ := h.MutateQuery(ctx, backend.DataQuery{
+			JSON: []byte(`{}`),
+		})
+
+		assert.NotEqual(t, ctx, newCtx)
+	})
+
+	t.Run("no grafana headers in context still works", func(t *testing.T) {
+		ctx := t.Context()
+
+		newCtx, _ := h.MutateQuery(ctx, backend.DataQuery{
+			JSON: []byte(`{}`),
+		})
+
+		assert.NotNil(t, newCtx)
+		_, ok := newCtx.Value(grafanaHeadersKey).(grafanaHeaders)
+		assert.False(t, ok)
+	})
+
+	t.Run("handles invalid JSON gracefully", func(t *testing.T) {
+		ctx := context.WithValue(t.Context(), grafanaHeadersKey, grafanaHeaders{
+			DashboardUID: "dash1",
+		})
+
+		newCtx, _ := h.MutateQuery(ctx, backend.DataQuery{
+			JSON: []byte(`invalid json`),
+		})
+
+		assert.NotEqual(t, ctx, newCtx)
+	})
+}