From 4fdaaf748a67cfdd7ee9da2968ebc6b3ce30abcc Mon Sep 17 00:00:00 2001 From: Charlie Le <3375195+CharlieTLe@users.noreply.github.com> Date: Wed, 24 Apr 2024 09:05:10 -0700 Subject: [PATCH] Set tenant id in prom analyse command (#31) * Set tenant id in prom analyse command - refactor roundtripper to be used in query_runner too - fixes https://github.com/cortexproject/cortex-tools/issues/13 Signed-off-by: Charlie Le * Update tests, changelog, docs Signed-off-by: Charlie Le --------- Signed-off-by: Charlie Le --- CHANGELOG.md | 1 + pkg/bench/query_runner.go | 21 ++++----------- pkg/commands/analyse_prometheus.go | 8 ++++-- pkg/httpmiddleware/roundtripper.go | 22 +++++++++++++++ pkg/httpmiddleware/roundtripper_test.go | 36 +++++++++++++++++++++++++ 5 files changed, 70 insertions(+), 18 deletions(-) create mode 100644 pkg/httpmiddleware/roundtripper.go create mode 100644 pkg/httpmiddleware/roundtripper_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index 36a4f50c1..095fe5b00 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ Order should be `CHANGE`, `FEATURE`, `ENHANCEMENT`, and `BUGFIX` * [CHANGE] Updates version of Go to 1.22 and Alpine to 3.19.1 in Dockerfiles * [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 ec42d5835..e75f785fb 100644 --- a/pkg/commands/analyse_prometheus.go +++ b/pkg/commands/analyse_prometheus.go @@ -17,6 +17,7 @@ import ( "gopkg.in/alecthomas/kingpin.v2" "github.com/cortexproject/cortex-tools/pkg/analyse" + "github.com/cortexproject/cortex-tools/pkg/httpmiddleware" ) type PrometheusAnalyseCommand struct { @@ -71,8 +72,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")) +}