Skip to content
This repository was archived by the owner on Mar 18, 2025. It is now read-only.

Commit 52c4189

Browse files
johncmingjohncming
authored
Add http headers support for write api. (#7)
* Add http headers support for write api. Signed-off-by: johncming <[email protected]> * Add http headers support for write api. * Add http headers support for write api.
1 parent 591f19a commit 52c4189

File tree

2 files changed

+154
-0
lines changed

2 files changed

+154
-0
lines changed

‎pkg/remotewrite/config.go

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import (
44
"encoding/json"
55
"net/url"
66
"strconv"
7+
"strings"
78
"time"
89

910
"github.com/kubernetes/helm/pkg/strvals"
@@ -25,6 +26,8 @@ type Config struct {
2526

2627
Url null.String `json:"url" envconfig:"K6_PROMETHEUS_REMOTE_URL"` // here, in the name of env variable, we assume that we won't need to distinguish between remote write URL vs remote read URL
2728

29+
Headers map[string]string `json:"headers" envconfig:"K6_PROMETHEUS_HEADERS"`
30+
2831
InsecureSkipTLSVerify null.Bool `json:"insecureSkipTLSVerify" envconfig:"K6_PROMETHEUS_INSECURE_SKIP_TLS_VERIFY"`
2932
CACert null.String `json:"caCertFile" envconfig:"K6_CA_CERT_FILE"`
3033

@@ -48,6 +51,7 @@ func NewConfig() Config {
4851
FlushPeriod: types.NullDurationFrom(defaultFlushPeriod),
4952
KeepTags: null.BoolFrom(true),
5053
KeepNameTag: null.BoolFrom(false),
54+
Headers: make(map[string]string),
5155
}
5256
}
5357

@@ -83,6 +87,7 @@ func (conf Config) ConstructRemoteConfig() (*remote.ClientConfig, error) {
8387
Timeout: model.Duration(defaultPrometheusTimeout),
8488
HTTPClientConfig: httpConfig,
8589
RetryOnRateLimit: true,
90+
Headers: conf.Headers,
8691
}
8792
return &remoteConfig, nil
8893
}
@@ -126,6 +131,12 @@ func (base Config) Apply(applied Config) Config {
126131
base.KeepNameTag = applied.KeepNameTag
127132
}
128133

134+
if len(applied.Headers) > 0 {
135+
for k, v := range applied.Headers {
136+
base.Headers[k] = v
137+
}
138+
}
139+
129140
return base
130141
}
131142

@@ -175,6 +186,15 @@ func ParseArg(arg string) (Config, error) {
175186
c.KeepNameTag = null.BoolFrom(v)
176187
}
177188

189+
c.Headers = make(map[string]string)
190+
if v, ok := params["headers"].(map[string]interface{}); ok {
191+
for k, v := range v {
192+
if v, ok := v.(string); ok {
193+
c.Headers[k] = v
194+
}
195+
}
196+
}
197+
178198
return c, nil
179199
}
180200

@@ -201,6 +221,17 @@ func GetConsolidatedConfig(jsonRawConf json.RawMessage, env map[string]string, a
201221
return null.NewBool(false, false), nil
202222
}
203223

224+
getEnvMap := func(env map[string]string, prefix string) map[string]string {
225+
result := make(map[string]string)
226+
for ek, ev := range env {
227+
if strings.HasPrefix(ek, prefix) {
228+
k := strings.TrimPrefix(ek, prefix)
229+
result[k] = ev
230+
}
231+
}
232+
return result
233+
}
234+
204235
// envconfig is not processing some undefined vars (at least duration) so apply them manually
205236
if flushPeriod, flushPeriodDefined := env["K6_PROMETHEUS_FLUSH_PERIOD"]; flushPeriodDefined {
206237
if err := result.FlushPeriod.UnmarshalText([]byte(flushPeriod)); err != nil {
@@ -253,6 +284,11 @@ func GetConsolidatedConfig(jsonRawConf json.RawMessage, env map[string]string, a
253284
}
254285
}
255286

287+
envHeaders := getEnvMap(env, "K6_PROMETHEUS_HEADERS_")
288+
for k, v := range envHeaders {
289+
result.Headers[k] = v
290+
}
291+
256292
if arg != "" {
257293
argConf, err := ParseArg(arg)
258294
if err != nil {

‎pkg/remotewrite/config_test.go

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@ func TestApply(t *testing.T) {
2525
User: null.StringFrom("user"),
2626
Password: null.StringFrom("pass"),
2727
FlushPeriod: types.NullDurationFrom(10 * time.Second),
28+
Headers: map[string]string{
29+
"X-Header": "value",
30+
},
2831
}
2932

3033
// Defaults should be overwritten by valid values
@@ -36,6 +39,7 @@ func TestApply(t *testing.T) {
3639
assert.Equal(t, fullConfig.User, c.User)
3740
assert.Equal(t, fullConfig.Password, c.Password)
3841
assert.Equal(t, fullConfig.FlushPeriod, c.FlushPeriod)
42+
assert.Equal(t, fullConfig.Headers, c.Headers)
3943

4044
// Defaults shouldn't be impacted by invalid values
4145
c = NewConfig()
@@ -78,6 +82,11 @@ func TestConfigParseArg(t *testing.T) {
7882
assert.Nil(t, err)
7983
assert.Equal(t, null.StringFrom("http://prometheus.remote:3412/write"), c.Url)
8084
assert.Equal(t, types.NullDurationFrom(time.Second*2), c.FlushPeriod)
85+
86+
c, err = ParseArg("url=http://prometheus.remote:3412/write,headers.X-Header=value")
87+
assert.Nil(t, err)
88+
assert.Equal(t, null.StringFrom("http://prometheus.remote:3412/write"), c.Url)
89+
assert.Equal(t, map[string]string{"X-Header": "value"}, c.Headers)
8190
}
8291

8392
// testing both GetConsolidatedConfig and ConstructRemoteConfig here until it's future config refactor takes shape (k6 #883)
@@ -108,6 +117,7 @@ func TestConstructRemoteConfig(t *testing.T) {
108117
FlushPeriod: types.NullDurationFrom(defaultFlushPeriod),
109118
KeepTags: null.BoolFrom(true),
110119
KeepNameTag: null.BoolFrom(false),
120+
Headers: make(map[string]string),
111121
},
112122
errString: "",
113123
remoteConfig: &remote.ClientConfig{
@@ -136,6 +146,7 @@ func TestConstructRemoteConfig(t *testing.T) {
136146
FlushPeriod: types.NullDurationFrom(defaultFlushPeriod),
137147
KeepTags: null.BoolFrom(true),
138148
KeepNameTag: null.BoolFrom(false),
149+
Headers: make(map[string]string),
139150
},
140151
errString: "",
141152
remoteConfig: &remote.ClientConfig{
@@ -169,6 +180,112 @@ func TestConstructRemoteConfig(t *testing.T) {
169180
errString: "strconv.ParseBool",
170181
remoteConfig: nil,
171182
},
183+
"remote_write_with_headers_json": {
184+
jsonRaw: json.RawMessage(fmt.Sprintf(`{"url":"%s","mapping":"mapping", "headers":{"X-Header":"value"}}`, u.String())),
185+
env: nil,
186+
arg: "",
187+
config: Config{
188+
Mapping: null.NewString("mapping", true),
189+
Url: null.StringFrom(u.String()),
190+
InsecureSkipTLSVerify: null.BoolFrom(true),
191+
CACert: null.NewString("", false),
192+
User: null.NewString("", false),
193+
Password: null.NewString("", false),
194+
FlushPeriod: types.NullDurationFrom(defaultFlushPeriod),
195+
KeepTags: null.BoolFrom(true),
196+
KeepNameTag: null.BoolFrom(false),
197+
Headers: map[string]string{
198+
"X-Header": "value",
199+
},
200+
},
201+
errString: "",
202+
remoteConfig: &remote.ClientConfig{
203+
URL: &promConfig.URL{URL: u},
204+
Timeout: model.Duration(defaultPrometheusTimeout),
205+
HTTPClientConfig: promConfig.HTTPClientConfig{
206+
FollowRedirects: true,
207+
TLSConfig: promConfig.TLSConfig{
208+
InsecureSkipVerify: false,
209+
},
210+
},
211+
RetryOnRateLimit: false,
212+
Headers: map[string]string{
213+
"X-Header": "value",
214+
},
215+
},
216+
},
217+
"remote_write_with_headers_env": {
218+
jsonRaw: json.RawMessage(fmt.Sprintf(`{"url":"%s","mapping":"mapping", "headers":{"X-Header":"value"}}`, u.String())),
219+
env: map[string]string{
220+
"K6_PROMETHEUS_HEADERS_X-Header": "value_from_env",
221+
},
222+
arg: "",
223+
config: Config{
224+
Mapping: null.NewString("mapping", true),
225+
Url: null.StringFrom(u.String()),
226+
InsecureSkipTLSVerify: null.BoolFrom(true),
227+
CACert: null.NewString("", false),
228+
User: null.NewString("", false),
229+
Password: null.NewString("", false),
230+
FlushPeriod: types.NullDurationFrom(defaultFlushPeriod),
231+
KeepTags: null.BoolFrom(true),
232+
KeepNameTag: null.BoolFrom(false),
233+
Headers: map[string]string{
234+
"X-Header": "value_from_env",
235+
},
236+
},
237+
errString: "",
238+
remoteConfig: &remote.ClientConfig{
239+
URL: &promConfig.URL{URL: u},
240+
Timeout: model.Duration(defaultPrometheusTimeout),
241+
HTTPClientConfig: promConfig.HTTPClientConfig{
242+
FollowRedirects: true,
243+
TLSConfig: promConfig.TLSConfig{
244+
InsecureSkipVerify: false,
245+
},
246+
},
247+
RetryOnRateLimit: false,
248+
Headers: map[string]string{
249+
"X-Header": "value_from_env",
250+
},
251+
},
252+
},
253+
"remote_write_with_headers_arg": {
254+
jsonRaw: json.RawMessage(fmt.Sprintf(`{"url":"%s","mapping":"mapping", "headers":{"X-Header":"value"}}`, u.String())),
255+
env: map[string]string{
256+
"K6_PROMETHEUS_HEADERS_X-Header": "value_from_env",
257+
},
258+
arg: "headers.X-Header=value_from_arg",
259+
config: Config{
260+
Mapping: null.NewString("mapping", true),
261+
Url: null.StringFrom(u.String()),
262+
InsecureSkipTLSVerify: null.BoolFrom(true),
263+
CACert: null.NewString("", false),
264+
User: null.NewString("", false),
265+
Password: null.NewString("", false),
266+
FlushPeriod: types.NullDurationFrom(defaultFlushPeriod),
267+
KeepTags: null.BoolFrom(true),
268+
KeepNameTag: null.BoolFrom(false),
269+
Headers: map[string]string{
270+
"X-Header": "value_from_arg",
271+
},
272+
},
273+
errString: "",
274+
remoteConfig: &remote.ClientConfig{
275+
URL: &promConfig.URL{URL: u},
276+
Timeout: model.Duration(defaultPrometheusTimeout),
277+
HTTPClientConfig: promConfig.HTTPClientConfig{
278+
FollowRedirects: true,
279+
TLSConfig: promConfig.TLSConfig{
280+
InsecureSkipVerify: false,
281+
},
282+
},
283+
RetryOnRateLimit: false,
284+
Headers: map[string]string{
285+
"X-Header": "value_from_arg",
286+
},
287+
},
288+
},
172289
}
173290

174291
for name, testCase := range testCases {
@@ -198,6 +315,7 @@ func assertConfig(t *testing.T, actual, expected Config) {
198315
assert.Equal(t, expected.FlushPeriod, actual.FlushPeriod)
199316
assert.Equal(t, expected.KeepTags, actual.KeepTags)
200317
assert.Equal(t, expected.KeepNameTag, expected.KeepNameTag)
318+
assert.Equal(t, expected.Headers, actual.Headers)
201319
}
202320

203321
func assertRemoteConfig(t *testing.T, actual, expected *remote.ClientConfig) {

0 commit comments

Comments
 (0)