From 688d7aac0f11cb8a1a2567a2c336664d61eaa797 Mon Sep 17 00:00:00 2001
From: addshore <addshorewiki@gmail.com>
Date: Mon, 19 Aug 2024 14:31:53 +0100
Subject: [PATCH] Set up to 10 custom headers on go client

Fixes #483
---
 pkg/client.go                   |  6 ++++
 pkg/datasource_settings.go      | 47 +++++++++++++++++++++++++-
 pkg/datasource_settings_test.go | 59 +++++++++++++++++++++++++++++++++
 3 files changed, 111 insertions(+), 1 deletion(-)
 create mode 100644 pkg/datasource_settings_test.go

diff --git a/pkg/client.go b/pkg/client.go
index b23420448..abf23455b 100644
--- a/pkg/client.go
+++ b/pkg/client.go
@@ -16,6 +16,7 @@ import (
 
 	"compress/flate"
 	"compress/gzip"
+
 	"github.com/andybalholm/brotli"
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/klauspost/compress/zstd"
@@ -76,6 +77,11 @@ func (client *ClickHouseClient) Query(ctx context.Context, query string) (*Respo
 			req.Header.Set("X-ClickHouse-Key", password)
 		}
 	}
+	if client.settings.CustomHeaders != nil {
+		for k, v := range client.settings.CustomHeaders {
+			req.Header.Set(k, v)
+		}
+	}
 
 	tlsCACert, tlsCACertExists := client.settings.Instance.DecryptedSecureJSONData["tlsCACert"]
 	tlsClientCert, tlsClientCertExists := client.settings.Instance.DecryptedSecureJSONData["tlsClientCert"]
diff --git a/pkg/datasource_settings.go b/pkg/datasource_settings.go
index 48cb5c393..58129a7a2 100644
--- a/pkg/datasource_settings.go
+++ b/pkg/datasource_settings.go
@@ -9,7 +9,6 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
 )
 
-// DatasourceSettings TODO Add Support custom headers
 type DatasourceSettings struct {
 	Instance backend.DataSourceInstanceSettings
 
@@ -22,6 +21,18 @@ type DatasourceSettings struct {
 	UseCompression              bool   `json:"useCompression,omitempty"`
 	CompressionType             string `json:"compressionType,omitempty"`
 	TLSSkipVerify               bool   `json:"tlsSkipVerify"`
+
+	CustomHeaders    map[string]string
+	HttpHeaderName1  string `json:"httpHeaderName1,omitempty"`
+	HttpHeaderName2  string `json:"httpHeaderName2,omitempty"`
+	HttpHeaderName3  string `json:"httpHeaderName3,omitempty"`
+	HttpHeaderName4  string `json:"httpHeaderName4,omitempty"`
+	HttpHeaderName5  string `json:"httpHeaderName5,omitempty"`
+	HttpHeaderName6  string `json:"httpHeaderName6,omitempty"`
+	HttpHeaderName7  string `json:"httpHeaderName7,omitempty"`
+	HttpHeaderName8  string `json:"httpHeaderName8,omitempty"`
+	HttpHeaderName9  string `json:"httpHeaderName9,omitempty"`
+	HttpHeaderName10 string `json:"httpHeaderName10,omitempty"`
 }
 
 func NewDatasourceSettings(ctx context.Context, settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
@@ -32,6 +43,40 @@ func NewDatasourceSettings(ctx context.Context, settings backend.DataSourceInsta
 		return nil, fmt.Errorf("unable to parse settings json %s. Error: %w", settings.JSONData, err)
 	}
 
+	dsSettings.CustomHeaders = make(map[string]string)
+	for i := 1; i <= 11; i++ {
+		headerName := ""
+		switch i {
+		case 1:
+			headerName = dsSettings.HttpHeaderName1
+		case 2:
+			headerName = dsSettings.HttpHeaderName2
+		case 3:
+			headerName = dsSettings.HttpHeaderName3
+		case 4:
+			headerName = dsSettings.HttpHeaderName4
+		case 5:
+			headerName = dsSettings.HttpHeaderName5
+		case 6:
+			headerName = dsSettings.HttpHeaderName6
+		case 7:
+			headerName = dsSettings.HttpHeaderName7
+		case 8:
+			headerName = dsSettings.HttpHeaderName8
+		case 9:
+			headerName = dsSettings.HttpHeaderName9
+		case 10:
+			headerName = dsSettings.HttpHeaderName10
+		case 11:
+			return nil, fmt.Errorf("too many custom headers")
+		}
+		if headerName == "" {
+			break
+		}
+		headerValue := settings.DecryptedSecureJSONData[fmt.Sprintf("httpHeaderValue%d", i)]
+		dsSettings.CustomHeaders[headerName] = headerValue
+	}
+
 	dsSettings.Instance = settings
 
 	return &dsSettings, nil
diff --git a/pkg/datasource_settings_test.go b/pkg/datasource_settings_test.go
new file mode 100644
index 000000000..0d958b403
--- /dev/null
+++ b/pkg/datasource_settings_test.go
@@ -0,0 +1,59 @@
+package main
+
+import (
+	"context"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewDatasourceSettings(t *testing.T) {
+	ctx := context.Background()
+
+	// Create a mock backend.DataSourceInstanceSettings
+	settings := backend.DataSourceInstanceSettings{
+		JSONData: []byte(`{
+			"addCorsHeader": true,
+			"defaultDatabase": "myDatabase",
+			"usePOST": false,
+			"useYandexCloudAuthorization": true,
+			"xHeaderKey": "myKey",
+			"xHeaderUser": "myUser",
+			"useCompression": true,
+			"compressionType": "gzip",
+			"tlsSkipVerify": false,
+			"httpHeaderName1": "header1",
+			"httpHeaderValue1": "value1",
+			"httpHeaderName2": "header2",
+			"httpHeaderValue2": "value2"
+		}`),
+		DecryptedSecureJSONData: map[string]string{
+			"httpHeaderValue1": "value1",
+			"httpHeaderValue2": "value2",
+		},
+	}
+
+	// Call the NewDatasourceSettings function
+	instance, err := NewDatasourceSettings(ctx, settings)
+	require.NoError(t, err)
+
+	// Assert the instance type
+	dsSettings, ok := instance.(*DatasourceSettings)
+	require.True(t, ok)
+
+	// Assert the instance fields
+	require.Equal(t, true, dsSettings.AddCorsHeader)
+	require.Equal(t, "myDatabase", dsSettings.DefaultDatabase)
+	require.Equal(t, false, dsSettings.UsePost)
+	require.Equal(t, true, dsSettings.UseYandexCloudAuthorization)
+	require.Equal(t, "myKey", dsSettings.XHeaderKey)
+	require.Equal(t, "myUser", dsSettings.XHeaderUser)
+	require.Equal(t, true, dsSettings.UseCompression)
+	require.Equal(t, "gzip", dsSettings.CompressionType)
+	require.Equal(t, false, dsSettings.TLSSkipVerify)
+	require.Equal(t, "header1", dsSettings.HttpHeaderName1)
+	require.Equal(t, "value1", dsSettings.CustomHeaders["header1"])
+	require.Equal(t, "header2", dsSettings.HttpHeaderName2)
+	require.Equal(t, "value2", dsSettings.CustomHeaders["header2"])
+}