diff --git a/pkg/bench/query_runner.go b/pkg/bench/query_runner.go index 362e1a81c..20a3eabd4 100644 --- a/pkg/bench/query_runner.go +++ b/pkg/bench/query_runner.go @@ -4,7 +4,6 @@ import ( "context" "flag" "math/rand" - "net/http" "sync" "time" @@ -19,6 +18,8 @@ import ( config_util "github.com/prometheus/common/config" "github.com/thanos-io/thanos/pkg/discovery/dns" "github.com/thanos-io/thanos/pkg/extprom" + + "github.com/cortexproject/cortex-tools/pkg/httpmiddleware" ) type QueryConfig struct { @@ -139,24 +140,12 @@ func (q *queryRunner) queryWorker(queryChan chan query) { } } -type tenantIDRoundTripper struct { - tenantName string - next http.RoundTripper -} - -func (r *tenantIDRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { - if r.tenantName != "" { - req.Header.Set("X-Scope-OrgID", r.tenantName) - } - return r.next.RoundTrip(req) -} - func newQueryClient(url, tenantName, username, password string) (v1.API, error) { apiClient, err := api.NewClient(api.Config{ Address: url, - RoundTripper: &tenantIDRoundTripper{ - tenantName: tenantName, - next: config_util.NewBasicAuthRoundTripper(username, config_util.Secret(password), "", api.DefaultRoundTripper), + RoundTripper: &httpmiddleware.TenantIDRoundTripper{ + TenantName: tenantName, + Next: config_util.NewBasicAuthRoundTripper(username, config_util.Secret(password), "", api.DefaultRoundTripper), }, }) diff --git a/pkg/commands/analyse_prometheus.go b/pkg/commands/analyse_prometheus.go index b5fa42ca6..068bfe842 100644 --- a/pkg/commands/analyse_prometheus.go +++ b/pkg/commands/analyse_prometheus.go @@ -18,6 +18,7 @@ import ( "gopkg.in/alecthomas/kingpin.v2" "github.com/cortexproject/cortex-tools/pkg/analyse" + "github.com/cortexproject/cortex-tools/pkg/httpmiddleware" ) type PrometheusAnalyseCommand struct { @@ -72,8 +73,11 @@ func (cmd *PrometheusAnalyseCommand) run(_ *kingpin.ParseContext) error { rt = config.NewBasicAuthRoundTripper(cmd.username, config.Secret(cmd.password), "", api.DefaultRoundTripper) } promClient, err := api.NewClient(api.Config{ - Address: cmd.address, - RoundTripper: rt, + Address: cmd.address, + RoundTripper: &httpmiddleware.TenantIDRoundTripper{ + TenantName: cmd.username, + Next: rt, + }, }) if err != nil { return err diff --git a/pkg/httpmiddleware/roundtripper.go b/pkg/httpmiddleware/roundtripper.go new file mode 100644 index 000000000..3a2770b50 --- /dev/null +++ b/pkg/httpmiddleware/roundtripper.go @@ -0,0 +1,15 @@ +package httpmiddleware + +import "net/http" + +type TenantIDRoundTripper struct { + TenantName string + Next http.RoundTripper +} + +func (r *TenantIDRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + if r.TenantName != "" { + req.Header.Set("X-Scope-OrgID", r.TenantName) + } + return r.Next.RoundTrip(req) +}