Allow injection of arbitrary headers in promtool, for auth etc. (#4389)

Signed-off-by: Tom Wilkie <tom.wilkie@gmail.com>
This commit is contained in:
Tom Wilkie 2019-06-30 11:50:23 +01:00 committed by GitHub
parent 32c3a1beef
commit 851131b074
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 16 additions and 5 deletions

View file

@ -19,6 +19,7 @@ import (
"encoding/json"
"fmt"
"math"
"net/http"
"net/url"
"os"
"path/filepath"
@ -30,6 +31,7 @@ import (
"github.com/pkg/errors"
"github.com/prometheus/client_golang/api"
v1 "github.com/prometheus/client_golang/api/prometheus/v1"
"github.com/prometheus/client_golang/prometheus/promhttp"
config_util "github.com/prometheus/common/config"
"github.com/prometheus/common/model"
"github.com/prometheus/common/version"
@ -70,6 +72,7 @@ func main() {
queryRangeCmd := queryCmd.Command("range", "Run range query.")
queryRangeServer := queryRangeCmd.Arg("server", "Prometheus server to query.").Required().String()
queryRangeExpr := queryRangeCmd.Arg("expr", "PromQL query expression.").Required().String()
queryRangeHeaders := queryRangeCmd.Flag("header", "Extra headers to send to server.").StringMap()
queryRangeBegin := queryRangeCmd.Flag("start", "Query range start time (RFC3339 or Unix timestamp).").String()
queryRangeEnd := queryRangeCmd.Flag("end", "Query range end time (RFC3339 or Unix timestamp).").String()
queryRangeStep := queryRangeCmd.Flag("step", "Query step size (duration).").Duration()
@ -123,7 +126,7 @@ func main() {
os.Exit(QueryInstant(*queryServer, *queryExpr, p))
case queryRangeCmd.FullCommand():
os.Exit(QueryRange(*queryRangeServer, *queryRangeExpr, *queryRangeBegin, *queryRangeEnd, *queryRangeStep, p))
os.Exit(QueryRange(*queryRangeServer, *queryRangeHeaders, *queryRangeExpr, *queryRangeBegin, *queryRangeEnd, *queryRangeStep, p))
case querySeriesCmd.FullCommand():
os.Exit(QuerySeries(*querySeriesServer, *querySeriesMatch, *querySeriesBegin, *querySeriesEnd, p))
@ -143,7 +146,6 @@ func main() {
case testRulesCmd.FullCommand():
os.Exit(RulesUnitTest(*testRulesFiles...))
}
}
// CheckConfig validates configuration files.
@ -361,11 +363,20 @@ func QueryInstant(url, query string, p printer) int {
}
// QueryRange performs a range query against a Prometheus server.
func QueryRange(url, query, start, end string, step time.Duration, p printer) int {
func QueryRange(url string, headers map[string]string, query, start, end string, step time.Duration, p printer) int {
config := api.Config{
Address: url,
}
if len(headers) > 0 {
config.RoundTripper = promhttp.RoundTripperFunc(func(req *http.Request) (*http.Response, error) {
for key, value := range headers {
req.Header.Add(key, value)
}
return http.DefaultTransport.RoundTrip(req)
})
}
// Create new client.
c, err := api.NewClient(config)
if err != nil {

View file

@ -26,7 +26,7 @@ func TestQueryRange(t *testing.T) {
defer s.Close()
p := &promqlPrinter{}
exitCode := QueryRange(s.URL, "up", "0", "300", 0, p)
exitCode := QueryRange(s.URL, map[string]string{}, "up", "0", "300", 0, p)
expectedPath := "/api/v1/query_range"
gotPath := getRequest().URL.Path
if gotPath != expectedPath {
@ -45,7 +45,7 @@ func TestQueryRange(t *testing.T) {
t.Error()
}
exitCode = QueryRange(s.URL, "up", "0", "300", 10*time.Millisecond, p)
exitCode = QueryRange(s.URL, map[string]string{}, "up", "0", "300", 10*time.Millisecond, p)
gotPath = getRequest().URL.Path
if gotPath != expectedPath {
t.Errorf("unexpected URL path %s (wanted %s)", gotPath, expectedPath)