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 #13 Signed-off-by: Charlie Le <[email protected]> * Update tests, changelog, docs Signed-off-by: Charlie Le <[email protected]> --------- Signed-off-by: Charlie Le <[email protected]>
cortexproject · May 9, 2024 · f89c799 · f89c799
1 parent 628b6ea
commit f89c799
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 18 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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"))
+}