diff --git a/pkg/plugin/driver.go b/pkg/plugin/driver.go index 60c7d318..0ed0ffa6 100644 --- a/pkg/plugin/driver.go +++ b/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 @@ -354,7 +364,20 @@ func (h *Clickhouse) Settings(ctx context.Context, config backend.DataSourceInst } } +// MutateQueryData extracts Grafana contextual headers from the request and +// stores them in the context for ClickHouse query metadata injection. +// It also preprocesses Grafana SQL queries. 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 +438,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, })) } diff --git a/pkg/plugin/driver_test.go b/pkg/plugin/driver_test.go index cc80bfd4..e86bbcf9 100644 --- a/pkg/plugin/driver_test.go +++ b/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) + }) +}