diff --git a/CHANGELOG.md b/CHANGELOG.md index a93c76377..0474806b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ Order should be `CHANGE`, `FEATURE`, `ENHANCEMENT`, and `BUGFIX` * [CHANGE] Upgrade to v1.13.2 * [FEATURE] Make rulerAPI Path configurable * [FEATURE] Add tool to deserialize alertmanager state file +* [BUGFIX] Set tenant id in prom analyse command ## v0.11.1 * [BUGFIX] Fix check for new version 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..8c95fa69b --- /dev/null +++ b/pkg/httpmiddleware/roundtripper.go @@ -0,0 +1,22 @@ +package httpmiddleware + +import "net/http" + +// TenantIDRoundTripper is a custom implementation of http.RoundTripper. +// It adds a tenant name to the request header before passing it to the next RoundTripper. +type TenantIDRoundTripper struct { + // TenantName is the name of the tenant to be added to the request header. + TenantName string + // Next is the next RoundTripper in the chain. + Next http.RoundTripper +} + +// RoundTrip adds the tenant name to the request header and then passes the request to the next RoundTripper. +// If TenantName is not set, it simply passes the request to the next RoundTripper. +// It returns the response from the next RoundTripper and any error encountered. +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) +} diff --git a/pkg/httpmiddleware/roundtripper_test.go b/pkg/httpmiddleware/roundtripper_test.go new file mode 100644 index 000000000..14bbda0a0 --- /dev/null +++ b/pkg/httpmiddleware/roundtripper_test.go @@ -0,0 +1,36 @@ +package httpmiddleware + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/stretchr/testify/require" +) + +// TestTenantIDRoundTripper_RoundTrip tests the RoundTrip method of TenantIDRoundTripper. +// It creates a new TenantIDRoundTripper with a TenantName and the default http transport as Next. +// It then sends a request and checks if the X-Scope-OrgID header of the request is correctly set to the TenantName. +// It also checks if the status code of the response is OK (200). +func TestTenantIDRoundTripper_RoundTrip(t *testing.T) { + // Create a new TenantIDRoundTripper with a TenantName and the default http transport as Next. + roundTripper := &TenantIDRoundTripper{ + TenantName: "test-tenant", + Next: http.DefaultTransport, + } + + // Create a new request. + req := httptest.NewRequest("GET", "https://example.com", nil) + + // Send the request using the RoundTrip method of the TenantIDRoundTripper. + resp, err := roundTripper.RoundTrip(req) + + // Check if there was an error. + require.NoError(t, err) + + // Check if the status code of the response is OK (200). + require.Equal(t, http.StatusOK, resp.StatusCode) + + // Check if the X-Scope-OrgID header of the request is correctly set to the TenantName. + require.Equal(t, "test-tenant", req.Header.Get("X-Scope-OrgID")) +}