Log HTTP headers as query_log.log_comment

pkg/plugin/driver.go

@@ -184,6 +184,15 @@ func (h *Clickhouse) Connect(ctx context.Context, config backend.DataSourceInsta
 		httpHeaders[k] = v
 	}
 
+	if settings.LogHeadersAsComment {
+		logComment, err := headersToLogComment(httpHeaders, settings.LogHeadersAsCommentRegex)
+		if err != nil {
+			backend.Logger.Warn("Failed to serialize headers as JSON", "error", err)
+		} else {
+			customSettings["log_comment"] = logComment
+		}
+	}
+
 	opts := &clickhouse.Options{
 		Addr: []string{fmt.Sprintf("%s:%d", settings.Host, settings.Port)},
 		Auth: clickhouse.Auth{
@@ -481,6 +490,22 @@ func extractForwardedHeadersFromMessage(message json.RawMessage) (map[string]str
 	return httpHeaders, nil
 }
 
+// headersToLogComment serializes the headers to a JSON string for use as a log comment.
+func headersToLogComment(headers map[string]string, regexPattern *regexp.Regexp) (string, error) {
+	// Compile the regex pattern
+	whiteListedHeaders := make(map[string]string)
+	for k, v := range headers {
+		if regexPattern.MatchString(k) {
+			whiteListedHeaders[k] = v
+		}
+	}
+	headersJSON, err := json.Marshal(whiteListedHeaders)
+	if err != nil {
+		return "", err
+	}
+	return string(headersJSON), nil
+}
+
 func mergeOpenTelemetryLabels(frame *data.Frame) error {
 	var attrFields []*data.Field
 	for _, field := range frame.Fields {

pkg/plugin/driver_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"regexp"
 	"testing"
 
 	"github.com/ClickHouse/clickhouse-go/v2"
@@ -317,3 +318,39 @@ func TestContainsClickHouseException(t *testing.T) {
 		assert.True(t, result)
 	})
 }
+
+func TestHeadersToLogComment(t *testing.T) {
+	regexPattern := "(?i)^(x-dashboard|x-panel|x-rule)"
+
+	t.Run("filters and serializes whitelisted headers", func(t *testing.T) {
+		headers := map[string]string{
+			"X-Dashboard-Id": "123",
+			"X-Panel-Id":     "456",
+			"X-Grafana-Id":   "should-be-excluded",
+			"Authorization":  "should-be-excluded",
+		}
+
+		result, err := headersToLogComment(headers, regexp.MustCompile(regexPattern))
+		assert.NoError(t, err)
+
+		var resultMap map[string]string
+		err = json.Unmarshal([]byte(result), &resultMap)
+		assert.NoError(t, err)
+
+		assert.Equal(t, "123", resultMap["X-Dashboard-Id"])
+		assert.Equal(t, "456", resultMap["X-Panel-Id"])
+		assert.NotContains(t, resultMap, "X-Grafana-Id")
+		assert.NotContains(t, resultMap, "Authorization")
+	})
+
+	t.Run("returns empty JSON object when no whitelisted headers", func(t *testing.T) {
+		headers := map[string]string{
+			"X-Grafana-Id":  "value",
+			"Authorization": "token",
+		}
+
+		result, err := headersToLogComment(headers, regexp.MustCompile(regexPattern))
+		assert.NoError(t, err)
+		assert.Equal(t, "{}", result)
+	})
+}

pkg/plugin/settings.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -40,10 +41,12 @@ type Settings struct {
 	MaxIdleConns    string `json:"maxIdleConns,omitempty"`
 	MaxOpenConns    string `json:"maxOpenConns,omitempty"`
 
-	HttpHeaders           map[string]string `json:"-"`
-	ForwardGrafanaHeaders bool              `json:"forwardGrafanaHeaders,omitempty"`
-	CustomSettings        []CustomSetting   `json:"customSettings"`
-	ProxyOptions          *proxy.Options
+	HttpHeaders              map[string]string `json:"-"`
+	ForwardGrafanaHeaders    bool              `json:"forwardGrafanaHeaders,omitempty"`
+	LogHeadersAsComment      bool              `json:"logHeadersAsComment,omitempty"`
+	LogHeadersAsCommentRegex *regexp.Regexp    `json:"logHeadersAsCommentRegex,omitempty"`
+	CustomSettings           []CustomSetting   `json:"customSettings"`
+	ProxyOptions             *proxy.Options
 
 	RowLimit       int64 `json:"rowLimit,omitempty"`
 	EnableRowLimit bool  `json:"enableRowLimit,omitempty"`
@@ -186,6 +189,27 @@ func LoadSettings(ctx context.Context, config backend.DataSourceInstanceSettings
 			settings.ForwardGrafanaHeaders = jsonData["forwardGrafanaHeaders"].(bool)
 		}
 	}
+	if jsonData["logHeadersAsComment"] != nil {
+		if logHeadersAsComment, ok := jsonData["logHeadersAsComment"].(string); ok {
+			settings.LogHeadersAsComment, err = strconv.ParseBool(logHeadersAsComment)
+			if err != nil {
+				return settings, backend.DownstreamError(fmt.Errorf("could not parse logHeadersAsComment value: %w", err))
+			}
+		} else {
+			settings.LogHeadersAsComment = jsonData["logHeadersAsComment"].(bool)
+		}
+	}
+	if jsonData["logHeadersAsCommentRegex"] != nil {
+		if regex, ok := jsonData["logHeadersAsCommentRegex"].(string); ok {
+			settings.LogHeadersAsCommentRegex, err = regexp.Compile(regex)
+			if err != nil {
+				return settings, backend.DownstreamError(fmt.Errorf("could not parse logHeadersAsCommentRegex value: %w", err))
+			}
+		}
+	}
+	if settings.LogHeadersAsComment && settings.LogHeadersAsCommentRegex == nil {
+		return settings, backend.DownstreamError(fmt.Errorf("logHeadersAsCommentRegex is required when logHeadersAsComment is true"))
+	}
 
 	if jsonData["enableRowLimit"] != nil {
 		if enableRowLimitString, ok := jsonData["enableRowLimit"].(string); ok {

src/components/configEditor/HttpHeadersConfig.test.tsx

@@ -15,6 +15,8 @@ describe('HttpHeadersConfig', () => {
         secureFields={{}}
         onHttpHeadersChange={() => {}}
         onForwardGrafanaHeadersChange={() => {}}
+        onLogHeadersAsCommentChange={() => {}}
+        onLogHeadersAsCommentRegexChange={() => {}}
       />
     );
     expect(result.container.firstChild).not.toBeNull();
@@ -29,6 +31,8 @@ describe('HttpHeadersConfig', () => {
         secureFields={{}}
         onHttpHeadersChange={onHttpHeadersChange}
         onForwardGrafanaHeadersChange={onForwardGrafanaHeadersChange}
+        onLogHeadersAsCommentChange={() => {}}
+        onLogHeadersAsCommentRegexChange={() => {}}
       />
     );
     expect(result.container.firstChild).not.toBeNull();
@@ -49,6 +53,8 @@ describe('HttpHeadersConfig', () => {
         secureFields={{}}
         onHttpHeadersChange={onHttpHeadersChange}
         onForwardGrafanaHeadersChange={onForwardGrafanaHeadersChange}
+        onLogHeadersAsCommentChange={() => {}}
+        onLogHeadersAsCommentRegexChange={() => {}}
       />
     );
     expect(result.container.firstChild).not.toBeNull();
@@ -98,6 +104,8 @@ describe('HttpHeadersConfig', () => {
         secureFields={{}}
         onHttpHeadersChange={onHttpHeadersChange}
         onForwardGrafanaHeadersChange={onForwardGrafanaHeadersChange}
+        onLogHeadersAsCommentChange={() => {}}
+        onLogHeadersAsCommentRegexChange={() => {}}
       />
     );
     expect(result.container.firstChild).not.toBeNull();
@@ -144,6 +152,8 @@ describe('forwardGrafanaHTTPHeaders', () => {
         secureFields={{}}
         onHttpHeadersChange={onHttpHeadersChange}
         onForwardGrafanaHeadersChange={onForwardGrafanaHeadersChange}
+        onLogHeadersAsCommentChange={() => {}}
+        onLogHeadersAsCommentRegexChange={() => {}}
       />
     );
     expect(result.container.firstChild).not.toBeNull();
@@ -154,3 +164,59 @@ describe('forwardGrafanaHTTPHeaders', () => {
     expect(onForwardGrafanaHeadersChange).toHaveBeenCalledTimes(1);
   });
 });
+
+describe('logHeadersAsComment', () => {
+  const selectors = allSelectors.components.Config.HttpHeaderConfig;
+
+  it('should call onLogHeadersAsCommentChange when switch is clicked', () => {
+    const onHttpHeadersChange = jest.fn();
+    const onForwardGrafanaHeadersChange = jest.fn();
+    const onLogHeadersAsCommentChange = jest.fn();
+    const result = render(
+      <HttpHeadersConfig
+        headers={[]}
+        secureFields={{}}
+        onHttpHeadersChange={onHttpHeadersChange}
+        onForwardGrafanaHeadersChange={onForwardGrafanaHeadersChange}
+        onLogHeadersAsCommentChange={onLogHeadersAsCommentChange}
+        onLogHeadersAsCommentRegexChange={() => {}}
+      />
+    );
+    expect(result.container.firstChild).not.toBeNull();
+
+    const logHeadersAsCommentSwitch = result.getByTestId(selectors.logHeadersAsCommentSwitch);
+    expect(logHeadersAsCommentSwitch).toBeInTheDocument();
+    fireEvent.click(logHeadersAsCommentSwitch);
+    expect(onLogHeadersAsCommentChange).toHaveBeenCalledTimes(1);
+    expect(onLogHeadersAsCommentChange).toHaveBeenCalledWith(true);
+  });
+
+  it('should call onLogHeadersAsCommentRegexChange when input is changed', () => {
+    const onHttpHeadersChange = jest.fn();
+    const onForwardGrafanaHeadersChange = jest.fn();
+    const onLogHeadersAsCommentChange = jest.fn();
+    const onLogHeadersAsCommentRegexChange = jest.fn();
+    const result = render(
+      <HttpHeadersConfig
+        headers={[]}
+        secureFields={{}}
+        logHeadersAsComment={true}
+        onHttpHeadersChange={onHttpHeadersChange}
+        onForwardGrafanaHeadersChange={onForwardGrafanaHeadersChange}
+        onLogHeadersAsCommentChange={onLogHeadersAsCommentChange}
+        onLogHeadersAsCommentRegexChange={onLogHeadersAsCommentRegexChange}
+      />
+    );
+    expect(result.container.firstChild).not.toBeNull();
+
+    const logHeadersAsCommentRegexInput = result.getByTestId(selectors.logHeadersAsCommentRegexInput);
+    expect(logHeadersAsCommentRegexInput).toBeInTheDocument();
+    fireEvent.change(logHeadersAsCommentRegexInput, { target: { value: 'test' } });
+    expect(onLogHeadersAsCommentRegexChange).toHaveBeenCalledTimes(1);
+    expect(onLogHeadersAsCommentRegexChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        target: expect.objectContaining({ value: 'test' }),
+      })
+    );
+  });
+});

src/components/configEditor/HttpHeadersConfig.tsx

@@ -10,16 +10,24 @@ import { KeyValue } from '@grafana/data';
 interface HttpHeadersConfigProps {
   headers?: CHHttpHeader[];
   forwardGrafanaHeaders?: boolean;
+  logHeadersAsComment?: boolean;
+  logHeadersAsCommentRegex?: string;
   secureFields: KeyValue<boolean>;
   onHttpHeadersChange: (v: CHHttpHeader[]) => void;
   onForwardGrafanaHeadersChange: (v: boolean) => void;
+  onLogHeadersAsCommentChange: (v: boolean) => void;
+  onLogHeadersAsCommentRegexChange: (e: ChangeEvent<HTMLInputElement>) => void;
 }
 
 export const HttpHeadersConfig = (props: HttpHeadersConfigProps) => {
   const { secureFields, onHttpHeadersChange } = props;
   const configuredSecureHeaders = useConfiguredSecureHttpHeaders(secureFields);
   const [headers, setHeaders] = useState<CHHttpHeader[]>(props.headers || []);
   const [forwardGrafanaHeaders, setForwardGrafanaHeaders] = useState<boolean>(props.forwardGrafanaHeaders || false);
+  const [logHeadersAsComment, setLogHeadersAsComment] = useState<boolean>(props.logHeadersAsComment || false);
+  const [logHeadersAsCommentRegex, setLogHeadersAsCommentRegex] = useState<string>(
+    props.logHeadersAsCommentRegex || allLabels.components.Config.HttpHeadersConfig.logHeadersAsCommentRegex.placeholder
+  );
   const labels = allLabels.components.Config.HttpHeadersConfig;
   const selectors = allSelectors.components.Config.HttpHeaderConfig;
 
@@ -41,6 +49,15 @@ export const HttpHeadersConfig = (props: HttpHeadersConfigProps) => {
     setForwardGrafanaHeaders(value);
     props.onForwardGrafanaHeadersChange(value);
   };
+  const updateLogHeadersAsComment = (value: boolean) => {
+    setLogHeadersAsComment(value);
+    props.onLogHeadersAsCommentChange(value);
+  };
+
+  const updateLogHeadersAsCommentRegex = (e: ChangeEvent<HTMLInputElement>) => {
+    setLogHeadersAsCommentRegex(e.target.value);
+    props.onLogHeadersAsCommentRegexChange(e);
+  };
 
   return (
     <ConfigSection title={labels.title}>
@@ -77,6 +94,24 @@ export const HttpHeadersConfig = (props: HttpHeadersConfigProps) => {
           onChange={(e) => updateForwardGrafanaHeaders(e.currentTarget.checked)}
         />
       </Field>
+      <Field label={labels.logHeadersAsComment.label} description={labels.logHeadersAsComment.tooltip}>
+        <Switch
+          data-testid={selectors.logHeadersAsCommentSwitch}
+          className={'gf-form'}
+          value={logHeadersAsComment}
+          onChange={(e) => updateLogHeadersAsComment(e.currentTarget.checked)}
+        />
+      </Field>
+      {logHeadersAsComment && (
+        <Field label={labels.logHeadersAsCommentRegex.label} description={labels.logHeadersAsCommentRegex.tooltip}>
+          <Input
+            data-testid={selectors.logHeadersAsCommentRegexInput}
+            value={logHeadersAsCommentRegex}
+            placeholder={labels.logHeadersAsCommentRegex.placeholder}
+            onChange={(e: ChangeEvent<HTMLInputElement>) => updateLogHeadersAsCommentRegex(e)}
+          />
+        </Field>
+      )}
     </ConfigSection>
   );
 };

src/labels.ts

@@ -91,6 +91,15 @@ export default {
           label: 'Forward Grafana HTTP Headers',
           tooltip: 'Forward Grafana HTTP Headers to datasource.',
         },
+        logHeadersAsComment: {
+          label: 'Log Headers as Comment',
+          tooltip: 'Serialize Grafana HTTP headers as JSON and include them as a log comment with each query.',
+        },
+        logHeadersAsCommentRegex: {
+          label: 'Header Whitelist Regex',
+          tooltip: `Regular expression to match headers that should be logged.`,
+          placeholder: '(?i)^(x-dashboard|x-panel|x-rule)',
+        },
       },
       AliasTableConfig: {
         title: 'Column Alias Tables',

src/selectors.ts

@@ -122,6 +122,8 @@ export const Components = {
       headerNameInput: 'config__http-header-config__header-name-input',
       headerValueInput: 'config__http-header-config__header-value-input',
       forwardGrafanaHeadersSwitch: 'config__http-header-config__forward-grafana-headers-switch',
+      logHeadersAsCommentSwitch: 'config__http-header-config__log-headers-as-comment-switch',
+      logHeadersAsCommentRegexInput: 'config__http-header-config__log-headers-as-comment-regex-input',
     },
     AliasTableConfig: {
       aliasEditor: 'config__alias-table-config__alias-editor',

src/types/config.ts

@@ -37,6 +37,8 @@ export interface CHConfig extends DataSourceJsonData {
 
   httpHeaders?: CHHttpHeader[];
   forwardGrafanaHeaders?: boolean;
+  logHeadersAsComment?: boolean;
+  logHeadersAsCommentRegex?: string;
 
   customSettings?: CHCustomSetting[];
   enableSecureSocksProxy?: boolean;