Răsfoiți Sursa

bearertoken

Ajay Tripathy 5 ani în urmă
părinte
comite
d97f4a1ca5
2 a modificat fișierele cu 32 adăugiri și 15 ștergeri
  1. 13 2
      pkg/env/costmodelenv.go
  2. 19 13
      pkg/prom/prom.go

+ 13 - 2
pkg/env/costmodelenv.go

@@ -28,10 +28,13 @@ const (
 	ErrorReportingEnabledEnvVar   = "ERROR_REPORTING_ENABLED"
 	ValuesReportingEnabledEnvVar  = "VALUES_REPORTING_ENABLED"
 
-	DBBasicAuthUsername           = "DB_BASIC_AUTH_USERNAME"
-	DBBasicAuthPassword           = "DB_BASIC_AUTH_PW"
+	DBBasicAuthUsername = "DB_BASIC_AUTH_USERNAME"
+	DBBasicAuthPassword = "DB_BASIC_AUTH_PW"
+	DBBearerToken       = "DB_BEARER_TOKEN"
+
 	MultiClusterBasicAuthUsername = "MC_BASIC_AUTH_USERNAME"
 	MultiClusterBasicAuthPassword = "MC_BASIC_AUTH_PW"
+	MultiClusterBearerToken       = "MC_BEARER_TOKEN"
 
 	InsecureSkipVerify = "INSECURE_SKIP_VERIFY"
 )
@@ -189,6 +192,10 @@ func GetDBBasicAuthUserPassword() string {
 
 }
 
+func GetDBBearerToken() string {
+	return Get(DBBearerToken, "")
+}
+
 // GetMultiClusterBasicAuthUsername returns the environemnt variable value for MultiClusterBasicAuthUsername
 func GetMultiClusterBasicAuthUsername() string {
 	return Get(MultiClusterBasicAuthUsername, "")
@@ -198,3 +205,7 @@ func GetMultiClusterBasicAuthUsername() string {
 func GetMultiClusterBasicAuthPassword() string {
 	return Get(MultiClusterBasicAuthPassword, "")
 }
+
+func GetMultiClusterToken() string {
+	return Get(MultiClusterBearerToken, "")
+}

+ 19 - 13
pkg/prom/prom.go

@@ -13,12 +13,13 @@ import (
 // Creates a new prometheus client which limits the total number of concurrent outbound requests
 // allowed at a given moment.
 type RateLimitedPrometheusClient struct {
-	client   prometheus.Client
-	limiter  *util.Semaphore
-	requests *util.AtomicInt32
-	outbound *util.AtomicInt32
-	username string
-	password string
+	client      prometheus.Client
+	limiter     *util.Semaphore
+	requests    *util.AtomicInt32
+	outbound    *util.AtomicInt32
+	username    string
+	password    string
+	bearerToken string
 }
 
 // requestCounter is used to determine if the prometheus client keeps track of
@@ -30,7 +31,7 @@ type requestCounter interface {
 
 // NewRateLimitedClient creates a prometheus client which limits the number of concurrent outbound
 // prometheus requests.
-func NewRateLimitedClient(config prometheus.Config, maxConcurrency int, username, password string) (prometheus.Client, error) {
+func NewRateLimitedClient(config prometheus.Config, maxConcurrency int, username, password, bearerToken string) (prometheus.Client, error) {
 	c, err := prometheus.NewClient(config)
 	if err != nil {
 		return nil, err
@@ -41,12 +42,13 @@ func NewRateLimitedClient(config prometheus.Config, maxConcurrency int, username
 	outbound := util.NewAtomicInt32(0)
 
 	return &RateLimitedPrometheusClient{
-		client:   c,
-		limiter:  limiter,
-		requests: requests,
-		outbound: outbound,
-		username: username,
-		password: password,
+		client:      c,
+		limiter:     limiter,
+		requests:    requests,
+		outbound:    outbound,
+		username:    username,
+		password:    password,
+		bearerToken: bearerToken,
 	}, nil
 }
 
@@ -84,6 +86,10 @@ func (rlpc *RateLimitedPrometheusClient) Do(ctx context.Context, req *http.Reque
 	if rlpc.username != "" {
 		req.SetBasicAuth(rlpc.username, rlpc.password)
 	}
+	if rlpc.bearerToken != "" {
+		token := "Bearer " + rlpc.bearerToken
+		req.Header.Add("Authorization", token)
+	}
 	// Increment the total request counter first
 	rlpc.requests.Increment()
 	defer rlpc.requests.Decrement()