app/vmctl: Add ability to set headers for vm-native HTTP requests. (#3906)

app/vmctl: Add ability to set headers for vm-native HTTP requests
This commit is contained in:
Gowtam Lal 2023-03-06 05:22:31 -05:00 committed by Aliaksandr Valialkin
parent e2527917a9
commit 0efdf47571
No known key found for this signature in database
GPG Key ID: A72BEC6CD3D0DED1
3 changed files with 72 additions and 0 deletions

View File

@ -328,10 +328,12 @@ const (
vmNativeSrcAddr = "vm-native-src-addr"
vmNativeSrcUser = "vm-native-src-user"
vmNativeSrcPassword = "vm-native-src-password"
vmNativeSrcHeaders = "vm-native-src-headers"
vmNativeDstAddr = "vm-native-dst-addr"
vmNativeDstUser = "vm-native-dst-user"
vmNativeDstPassword = "vm-native-dst-password"
vmNativeDstHeaders = "vm-native-dst-headers"
)
var (
@ -373,6 +375,12 @@ var (
Usage: "VictoriaMetrics password for basic auth",
EnvVars: []string{"VM_NATIVE_SRC_PASSWORD"},
},
&cli.StringFlag{
Name: vmNativeSrcHeaders,
Usage: "Optional HTTP headers to send with each request to the corresponding source address. \n" +
"For example, --vm-native-src-headers='My-Auth:foobar' would send 'My-Auth: foobar' HTTP header with every request to the corresponding source address. \n" +
"Multiple headers must be delimited by '^^': --vm-native-src-headers='header1:value1^^header2:value2'",
},
&cli.StringFlag{
Name: vmNativeDstAddr,
Usage: "VictoriaMetrics address to perform import to. \n" +
@ -390,6 +398,12 @@ var (
Usage: "VictoriaMetrics password for basic auth",
EnvVars: []string{"VM_NATIVE_DST_PASSWORD"},
},
&cli.StringFlag{
Name: vmNativeDstHeaders,
Usage: "Optional HTTP headers to send with each request to the corresponding destination address. \n" +
"For example, --vm-native-dst-headers='My-Auth:foobar' would send 'My-Auth: foobar' HTTP header with every request to the corresponding destination address. \n" +
"Multiple headers must be delimited by '^^': --vm-native-dst-headers='header1:value1^^header2:value2'",
},
&cli.StringSliceFlag{
Name: vmExtraLabel,
Value: nil,

View File

@ -212,12 +212,14 @@ func main() {
Addr: strings.Trim(c.String(vmNativeSrcAddr), "/"),
User: c.String(vmNativeSrcUser),
Password: c.String(vmNativeSrcPassword),
Headers: c.String(vmNativeSrcHeaders),
},
dst: &native.Client{
Addr: strings.Trim(c.String(vmNativeDstAddr), "/"),
User: c.String(vmNativeDstUser),
Password: c.String(vmNativeDstPassword),
ExtraLabels: c.StringSlice(vmExtraLabel),
Headers: c.String(vmNativeDstHeaders),
},
backoff: backoff.New(),
cc: c.Int(vmConcurrency),

View File

@ -6,6 +6,7 @@ import (
"fmt"
"io"
"net/http"
"strings"
)
const (
@ -21,6 +22,7 @@ type Client struct {
User string
Password string
ExtraLabels []string
Headers string
}
// LabelValues represents series from api/v1/series response
@ -89,6 +91,16 @@ func (c *Client) ImportPipe(ctx context.Context, dstURL string, pr *io.PipeReade
if err != nil {
return fmt.Errorf("cannot create import request to %q: %s", c.Addr, err)
}
parsedHeaders, err := parseHeaders(c.Headers)
if err != nil {
return err
}
for _, header := range parsedHeaders {
req.Header.Set(header.key, header.value)
}
importResp, err := c.do(req, http.StatusNoContent)
if err != nil {
return fmt.Errorf("import request failed: %s", err)
@ -118,6 +130,16 @@ func (c *Client) ExportPipe(ctx context.Context, url string, f Filter) (io.ReadC
// disable compression since it is meaningless for native format
req.Header.Set("Accept-Encoding", "identity")
parsedHeaders, err := parseHeaders(c.Headers)
if err != nil {
return nil, err
}
for _, header := range parsedHeaders {
req.Header.Set(header.key, header.value)
}
resp, err := c.do(req, http.StatusOK)
if err != nil {
return nil, fmt.Errorf("export request failed: %w", err)
@ -142,6 +164,15 @@ func (c *Client) GetSourceTenants(ctx context.Context, f Filter) ([]string, erro
}
req.URL.RawQuery = params.Encode()
parsedHeaders, err := parseHeaders(c.Headers)
if err != nil {
return nil, err
}
for _, header := range parsedHeaders {
req.Header.Set(header.key, header.value)
}
resp, err := c.do(req, http.StatusOK)
if err != nil {
return nil, fmt.Errorf("tenants request failed: %s", err)
@ -179,3 +210,28 @@ func (c *Client) do(req *http.Request, expSC int) (*http.Response, error) {
}
return resp, err
}
type keyValue struct {
key string
value string
}
func parseHeaders(headers string) ([]keyValue, error) {
if len(headers) == 0 {
return nil, nil
}
var headersSplitByDelimiter = strings.Split(headers, "^^")
kvs := make([]keyValue, len(headersSplitByDelimiter))
for i, h := range headersSplitByDelimiter {
n := strings.IndexByte(h, ':')
if n < 0 {
return nil, fmt.Errorf(`missing ':' in header %q; expecting "key: value" format`, h)
}
kv := &kvs[i]
kv.key = strings.TrimSpace(h[:n])
kv.value = strings.TrimSpace(h[n+1:])
}
return kvs, nil
}