datasource: add support for HTTP headers

Signed-off-by: Guillaume Delbergue <[email protected]>

Author: guyguy333

datasource.go

@@ -3,7 +3,10 @@ package gapi
 import (
 	"bytes"
 	"encoding/json"
+	"errors"
 	"fmt"
+	"regexp"
+	"strconv"
 )
 
 // DataSource represents a Grafana data source.
@@ -27,16 +30,52 @@ type DataSource struct {
 	// Deprecated: Use secureJsonData.basicAuthPassword instead.
 	BasicAuthPassword string `json:"basicAuthPassword,omitempty"`
 
+	// Helper to read/write http headers
+	HTTPHeaders map[string]string `json:"-"`
+
 	JSONData       JSONData       `json:"jsonData,omitempty"`
 	SecureJSONData SecureJSONData `json:"secureJsonData,omitempty"`
 }
 
+// Required to avoid recursion during (un)marshal
+type _DataSource DataSource
+
+// Marshal DataSource
+func (ds *DataSource) MarshalJSON() ([]byte, error) {
+	var index int64
+	dataSource := _DataSource(*ds)
+	dataSource.JSONData.httpHeaderNames = make(map[int64]string, len(ds.HTTPHeaders))
+	dataSource.SecureJSONData.httpHeaderValues = make(map[int64]string, len(ds.HTTPHeaders))
+	for name, value := range ds.HTTPHeaders {
+		dataSource.JSONData.httpHeaderNames[index] = name
+		dataSource.SecureJSONData.httpHeaderValues[index] = value
+		index++
+	}
+	return json.Marshal(dataSource)
+}
+
+// Unmarshal DataSource
+func (ds *DataSource) UnmarshalJSON(b []byte) (err error) {
+	dataSource := _DataSource(*ds)
+	if err = json.Unmarshal(b, &dataSource); err == nil {
+		*ds = DataSource(dataSource)
+	}
+	if len(ds.JSONData.httpHeaderNames) != len(ds.SecureJSONData.httpHeaderValues) {
+		return errors.New("HTTP headers names length doesn't match HTTP header values length")
+	}
+	for index, value := range ds.JSONData.httpHeaderNames {
+		ds.HTTPHeaders[value] = ds.SecureJSONData.httpHeaderValues[index]
+	}
+	return err
+}
+
 // JSONData is a representation of the datasource `jsonData` property
 type JSONData struct {
 	// Used by all datasources
 	TLSAuth           bool `json:"tlsAuth,omitempty"`
 	TLSAuthWithCACert bool `json:"tlsAuthWithCACert,omitempty"`
 	TLSSkipVerify     bool `json:"tlsSkipVerify,omitempty"`
+	httpHeaderNames   map[int64]string
 
 	// Used by Graphite
 	GraphiteVersion string `json:"graphiteVersion,omitempty"`
@@ -95,6 +134,49 @@ type JSONData struct {
 	SigV4Region        string `json:"sigV4Region,omitempty"`
 }
 
+// Required to avoid recursion during (un)marshal
+type _JSONData JSONData
+
+// Marshal JSONData
+func (jd JSONData) MarshalJSON() ([]byte, error) {
+	jsonData := _JSONData(jd)
+	b, err := json.Marshal(jsonData)
+	if err != nil {
+		return nil, err
+	}
+	fields := make(map[string]interface{})
+	if err = json.Unmarshal(b, &fields); err != nil {
+		return nil, err
+	}
+	for index, name := range jd.httpHeaderNames {
+		fields[fmt.Sprintf("httpHeaderName%d", index+1)] = name
+	}
+	return json.Marshal(fields)
+}
+
+// Unmarshal JSONData
+func (jd *JSONData) UnmarshalJSON(b []byte) (err error) {
+	jsonData := _JSONData(*jd)
+	if err = json.Unmarshal(b, &jsonData); err == nil {
+		*jd = JSONData(jsonData)
+	}
+	fields := make(map[string]interface{})
+	if err = json.Unmarshal(b, &fields); err == nil {
+		for name, value := range fields {
+			re := regexp.MustCompile("httpHeaderName([0-9]+)")
+			match := re.FindStringSubmatch(name)
+			if len(match) == 1 {
+				index, err := strconv.ParseInt(match[0], 10, 64)
+				if err != nil {
+					return err
+				}
+				jd.httpHeaderNames[index-1] = value.(string)
+			}
+		}
+	}
+	return err
+}
+
 // SecureJSONData is a representation of the datasource `secureJsonData` property
 type SecureJSONData struct {
 	// Used by all datasources
@@ -103,6 +185,7 @@ type SecureJSONData struct {
 	TLSClientKey      string `json:"tlsClientKey,omitempty"`
 	Password          string `json:"password,omitempty"`
 	BasicAuthPassword string `json:"basicAuthPassword,omitempty"`
+	httpHeaderValues  map[int64]string
 
 	// Used by Cloudwatch
 	AccessKey string `json:"accessKey,omitempty"`
@@ -116,6 +199,26 @@ type SecureJSONData struct {
 	SigV4SecretKey string `json:"sigV4SecretKey,omitempty"`
 }
 
+// Required to avoid recursion during unmarshal
+type _SecureJSONData SecureJSONData
+
+// Marshal SecureJSONData
+func (sjd SecureJSONData) MarshalJSON() ([]byte, error) {
+	secureJSONData := _SecureJSONData(sjd)
+	b, err := json.Marshal(secureJSONData)
+	if err != nil {
+		return nil, err
+	}
+	fields := make(map[string]interface{})
+	if err = json.Unmarshal(b, &fields); err != nil {
+		return nil, err
+	}
+	for index, value := range sjd.httpHeaderValues {
+		fields[fmt.Sprintf("httpHeaderValue%d", index+1)] = value
+	}
+	return json.Marshal(fields)
+}
+
 // NewDataSource creates a new Grafana data source.
 func (c *Client) NewDataSource(s *DataSource) (int64, error) {
 	data, err := json.Marshal(s)