Make remote write tenant header configurable (#25)

* make tenant header configurable Signed-off-by: Ian Billett <[email protected]> * rename round tripper Signed-off-by: Ian Billett <[email protected]> * formatting Signed-off-by: Ian Billett <[email protected]> * add health alias Signed-off-by: Ian Billett <[email protected]>
prometheus-community · Aug 2, 2022 · 19b624b · 19b624b
1 parent 54fd224
commit 19b624b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 8 deletions.
diff --git a/cmd/avalanche.go b/cmd/avalanche.go
@@ -48,6 +48,7 @@ var (
 	remoteReqsInterval  = kingpin.Flag("remote-write-interval", "delay between each remote write request.").Default("100ms").Duration()
 	remoteTenant        = kingpin.Flag("remote-tenant", "Tenant ID to include in remote_write send").Default("0").String()
 	tlsClientInsecure   = kingpin.Flag("tls-client-insecure", "Skip certificate check on tls connection").Default("false").Bool()
+	remoteTenantHeader  = kingpin.Flag("remote-tenant-header", "Tenant ID to include in remote_write send. The default, is the default tenant header expected by Cortex.").Default("X-Scope-OrgID").String()
 )
 
 func main() {
@@ -81,6 +82,7 @@ func main() {
 			TLSClientConfig: tls.Config{
 				InsecureSkipVerify: *tlsClientInsecure,
 			},
+			TenantHeader: *remoteTenantHeader,
 		}
 
 		// Collect Pprof during the write only if not collecting within a regular interval.

diff --git a/metrics/serve.go b/metrics/serve.go
@@ -21,7 +21,7 @@ import (
 	"sync"
 	"time"
 
-	"github.com/nelkinda/health-go"
+	health "github.com/nelkinda/health-go"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 )

diff --git a/metrics/write.go b/metrics/write.go
@@ -53,6 +53,7 @@ type ConfigWrite struct {
 	PprofURLs       []*url.URL
 	Tenant          string
 	TLSClientConfig tls.Config
+	TenantHeader    string
 }
 
 // Client for the remote write requests.
@@ -68,7 +69,7 @@ func SendRemoteWrite(config *ConfigWrite) error {
 	var rt http.RoundTripper = &http.Transport{
 		TLSClientConfig: &config.TLSClientConfig,
 	}
-	rt = &cortexTenantRoundTripper{tenant: config.Tenant, rt: rt}
+	rt = &tenantRoundTripper{tenant: config.Tenant, tenantHeader: config.TenantHeader, rt: rt}
 	httpClient := &http.Client{Transport: rt}
 
 	c := Client{
@@ -79,16 +80,17 @@ func SendRemoteWrite(config *ConfigWrite) error {
 	return c.write()
 }
 
-// Add the tenant ID header required by Cortex
-func (rt *cortexTenantRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+// Add the tenant ID header
+func (rt *tenantRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	req = cloneRequest(req)
-	req.Header.Set("X-Scope-OrgID", rt.tenant)
+	req.Header.Set(rt.tenantHeader, rt.tenant)
 	return rt.rt.RoundTrip(req)
 }
 
-type cortexTenantRoundTripper struct {
-	tenant string
-	rt     http.RoundTripper
+type tenantRoundTripper struct {
+	tenant       string
+	tenantHeader string
+	rt           http.RoundTripper
 }
 
 // cloneRequest returns a clone of the provided *http.Request.