Prometheus: Remove cache, pass headers in request, simplify client creation for resource calls and custom client

pkg/tsdb/prometheus/client/client.go

@@ -18,6 +18,9 @@ type doer interface {
 	Do(req *http.Request) (*http.Response, error)
 }
 
+// Client is a custom Prometheus client. Reason for this is that Prom Go client serializes response into its own
+// objects, we have to go through them and then serialize again into DataFrame which isn't very efficient. Using custom
+// client we can parse response directly into DataFrame.
 type Client struct {
 	doer    doer
 	method  string
@@ -28,101 +31,119 @@ func NewClient(d doer, method, baseUrl string) *Client {
 	return &Client{doer: d, method: method, baseUrl: baseUrl}
 }
 
-func (c *Client) QueryRange(ctx context.Context, q *models.Query) (*http.Response, error) {
-	u, err := url.ParseRequestURI(c.baseUrl)
+func (c *Client) QueryRange(ctx context.Context, q *models.Query, headers http.Header) (*http.Response, error) {
+	tr := q.TimeRange()
+	u, err := c.createUrl("api/v1/query_range", map[string]string{
+		"query": q.Expr,
+		"start": formatTime(tr.Start),
+		"end":   formatTime(tr.End),
+		"step":  strconv.FormatFloat(tr.Step.Seconds(), 'f', -1, 64),
+	})
 	if err != nil {
 		return nil, err
 	}
-
-	u.Path = path.Join(u.Path, "api/v1/query_range")
-
-	qs := u.Query()
-	qs.Set("query", q.Expr)
-	tr := q.TimeRange()
-	qs.Set("start", formatTime(tr.Start))
-	qs.Set("end", formatTime(tr.End))
-	qs.Set("step", strconv.FormatFloat(tr.Step.Seconds(), 'f', -1, 64))
-
-	return c.fetch(ctx, c.method, u, qs, nil)
-}
-
-func (c *Client) QueryInstant(ctx context.Context, q *models.Query) (*http.Response, error) {
-	u, err := url.ParseRequestURI(c.baseUrl)
+	req, err := createRequest(ctx, c.method, u, nil, headers)
 	if err != nil {
 		return nil, err
 	}
 
-	u.Path = path.Join(u.Path, "api/v1/query")
+	return c.doer.Do(req)
+}
 
-	qs := u.Query()
-	qs.Set("query", q.Expr)
+func (c *Client) QueryInstant(ctx context.Context, q *models.Query, headers http.Header) (*http.Response, error) {
+	qs := map[string]string{"query": q.Expr}
 	tr := q.TimeRange()
 	if !tr.End.IsZero() {
-		qs.Set("time", formatTime(tr.End))
+		qs["time"] = formatTime(tr.End)
 	}
 
-	return c.fetch(ctx, c.method, u, qs, nil)
-}
-
-func (c *Client) QueryExemplars(ctx context.Context, q *models.Query) (*http.Response, error) {
-	u, err := url.ParseRequestURI(c.baseUrl)
+	u, err := c.createUrl("api/v1/query", qs)
+	if err != nil {
+		return nil, err
+	}
+	req, err := createRequest(ctx, c.method, u, nil, headers)
 	if err != nil {
 		return nil, err
 	}
 
-	u.Path = path.Join(u.Path, "api/v1/query_exemplars")
+	return c.doer.Do(req)
+}
 
-	qs := u.Query()
+func (c *Client) QueryExemplars(ctx context.Context, q *models.Query, headers http.Header) (*http.Response, error) {
 	tr := q.TimeRange()
-	qs.Set("query", q.Expr)
-	qs.Set("start", formatTime(tr.Start))
-	qs.Set("end", formatTime(tr.End))
+	u, err := c.createUrl("api/v1/query_exemplars", map[string]string{
+		"query": q.Expr,
+		"start": formatTime(tr.Start),
+		"end":   formatTime(tr.End),
+	})
+	if err != nil {
+		return nil, err
+	}
 
-	return c.fetch(ctx, c.method, u, qs, nil)
-}
+	req, err := createRequest(ctx, c.method, u, nil, headers)
+	if err != nil {
+		return nil, err
+	}
 
-type FetchReq struct {
-	Method      string
-	Url         *url.URL
-	QueryString url.Values
+	return c.doer.Do(req)
 }
 
 func (c *Client) QueryResource(ctx context.Context, req *backend.CallResourceRequest) (*http.Response, error) {
 	// The way URL is represented in CallResourceRequest and what we need for the fetch function is different
 	// so here we have to do a bit of parsing, so we can then compose it with the base url in correct way.
-	baseUrlParsed, err := url.ParseRequestURI(c.baseUrl)
+	reqUrlParsed, err := url.Parse(req.URL)
 	if err != nil {
 		return nil, err
 	}
-	reqUrlParsed, err := url.Parse(req.URL)
+	u, err := c.createUrl(req.Path, nil)
 	if err != nil {
 		return nil, err
 	}
+	u.RawQuery = reqUrlParsed.RawQuery
 
-	baseUrlParsed.Path = path.Join(baseUrlParsed.Path, req.Path)
-	baseUrlParsed.RawQuery = reqUrlParsed.RawQuery
+	// We use method from the request, as for resources front end may do a fallback to GET if POST does not work
+	// nad we want to respect that.
+	httpRequest, err := createRequest(ctx, req.Method, u, req.Body, req.Headers)
+	if err != nil {
+		return nil, err
+	}
 
-	return c.fetch(ctx, req.Method, baseUrlParsed, nil, req.Body)
+	return c.doer.Do(httpRequest)
 }
 
-func (c *Client) fetch(ctx context.Context, method string, u *url.URL, qs url.Values, body []byte) (*http.Response, error) {
-	// The qs arg seems to be used in some callers of this method, but you can already pass them in the URL object
-	if strings.ToUpper(method) == http.MethodGet && qs != nil {
-		u.RawQuery = qs.Encode()
+func (c *Client) createUrl(endpoint string, qs map[string]string) (*url.URL, error) {
+	finalUrl, err := url.ParseRequestURI(c.baseUrl)
+	if err != nil {
+		return nil, err
+	}
+
+	finalUrl.Path = path.Join(finalUrl.Path, endpoint)
+	urlQuery := finalUrl.Query()
+
+	for key, val := range qs {
+		urlQuery.Set(key, val)
 	}
+
+	finalUrl.RawQuery = urlQuery.Encode()
+	return finalUrl, nil
+}
+
+func createRequest(ctx context.Context, method string, u *url.URL, body []byte, header http.Header) (*http.Request, error) {
 	bodyReader := bytes.NewReader(body)
 	request, err := http.NewRequestWithContext(ctx, method, u.String(), bodyReader)
 	if err != nil {
 		return nil, err
 	}
-
+	// request.Header is created empty from NewRequestWithContext so we can just replace it
+	if header != nil {
+		request.Header = header
+	}
 	// This may not be true but right now we don't have more information here and seems like we send just this type
 	// of encoding right now if it is a POST
 	if strings.ToUpper(method) == http.MethodPost {
 		request.Header.Set("Content-Type", "application/x-www-form-urlencoded")
 	}
-
-	return c.doer.Do(request)
+	return request, nil
 }
 
 func formatTime(t time.Time) string {