/
client.go
73 lines (61 loc) · 1.93 KB
/
client.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
package main
import (
"fmt"
"net/http"
"github.com/prometheus/common/version"
"gopkg.in/alecthomas/kingpin.v2"
)
const (
envPrefix = "PROFILECLI_"
)
var userAgentHeader = fmt.Sprintf("pyroscope/%s", version.Version)
type phlareClient struct {
TenantID string
URL string
BasicAuth struct {
Username string
Password string
}
defaultTransport http.RoundTripper
client *http.Client
}
type authRoundTripper struct {
client *phlareClient
next http.RoundTripper
}
func (a *authRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if c := a.client; c != nil {
if c.TenantID != "" {
req.Header.Set("X-Scope-OrgID", c.TenantID)
}
if c.BasicAuth.Username != "" || c.BasicAuth.Password != "" {
req.SetBasicAuth(c.BasicAuth.Username, c.BasicAuth.Password)
}
}
req.Header.Set("User-Agent", userAgentHeader)
return a.next.RoundTrip(req)
}
func (c *phlareClient) httpClient() *http.Client {
if c.client == nil {
if c.defaultTransport == nil {
c.defaultTransport = http.DefaultTransport
}
c.client = &http.Client{Transport: &authRoundTripper{
client: c,
next: c.defaultTransport,
}}
}
return c.client
}
type commander interface {
Flag(name, help string) *kingpin.FlagClause
Arg(name, help string) *kingpin.ArgClause
}
func addPhlareClient(cmd commander) *phlareClient {
client := &phlareClient{}
cmd.Flag("url", "URL of the profile store.").Default("http://localhost:4040").Envar(envPrefix + "URL").StringVar(&client.URL)
cmd.Flag("tenant-id", "The tenant ID to be used for the X-Scope-OrgID header.").Default("").Envar(envPrefix + "TENANT_ID").StringVar(&client.TenantID)
cmd.Flag("username", "The username to be used for basic auth.").Default("").Envar(envPrefix + "USERNAME").StringVar(&client.BasicAuth.Username)
cmd.Flag("password", "The password to be used for basic auth.").Default("").Envar(envPrefix + "PASSWORD").StringVar(&client.BasicAuth.Password)
return client
}