pass request with context on client call

grafana · Aug 8, 2023 · bf3e9a4 · bf3e9a4
1 parent 2f905b9
commit bf3e9a4
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 13 deletions.
diff --git a/pkg/client/alerts.go b/pkg/client/alerts.go
@@ -26,7 +26,7 @@ func (r *CortexClient) CreateAlertmanagerConfig(ctx context.Context, cfg string,
 		return err
 	}
 
-	res, err := r.doRequest(alertmanagerAPIPath, "POST", payload)
+	res, err := r.doRequest(ctx, alertmanagerAPIPath, "POST", payload)
 	if err != nil {
 		return err
 	}
@@ -38,7 +38,7 @@ func (r *CortexClient) CreateAlertmanagerConfig(ctx context.Context, cfg string,
 
 // DeleteAlermanagerConfig deletes the users alertmanagerconfig
 func (r *CortexClient) DeleteAlermanagerConfig(ctx context.Context) error {
-	res, err := r.doRequest(alertmanagerAPIPath, "DELETE", nil)
+	res, err := r.doRequest(ctx, alertmanagerAPIPath, "DELETE", nil)
 	if err != nil {
 		return err
 	}
@@ -50,7 +50,7 @@ func (r *CortexClient) DeleteAlermanagerConfig(ctx context.Context) error {
 
 // GetAlertmanagerConfig retrieves a rule group
 func (r *CortexClient) GetAlertmanagerConfig(ctx context.Context) (string, map[string]string, error) {
-	res, err := r.doRequest(alertmanagerAPIPath, "GET", nil)
+	res, err := r.doRequest(ctx, alertmanagerAPIPath, "GET", nil)
 	if err != nil {
 		log.Debugln("no alert config present in response")
 		return "", nil, err

diff --git a/pkg/client/client.go b/pkg/client/client.go
@@ -103,16 +103,16 @@ func (r *CortexClient) Query(ctx context.Context, query string) (*http.Response,
 	query = fmt.Sprintf("query=%s&time=%d", query, time.Now().Unix())
 	escapedQuery := url.PathEscape(query)
 
-	res, err := r.doRequest("/api/prom/api/v1/query?"+escapedQuery, "GET", nil)
+	res, err := r.doRequest(ctx, "/api/prom/api/v1/query?"+escapedQuery, "GET", nil)
 	if err != nil {
 		return nil, err
 	}
 
 	return res, nil
 }
 
-func (r *CortexClient) doRequest(path, method string, payload []byte) (*http.Response, error) {
-	req, err := buildRequest(path, method, *r.endpoint, payload)
+func (r *CortexClient) doRequest(ctx context.Context, path, method string, payload []byte) (*http.Response, error) {
+	req, err := buildRequest(ctx, path, method, *r.endpoint, payload)
 	if err != nil {
 		return nil, err
 	}
@@ -205,7 +205,7 @@ func joinPath(baseURLPath, targetPath string) string {
 	return strings.TrimSuffix(baseURLPath, "/") + targetPath
 }
 
-func buildRequest(p, m string, endpoint url.URL, payload []byte) (*http.Request, error) {
+func buildRequest(ctx context.Context, p, m string, endpoint url.URL, payload []byte) (*http.Request, error) {
 	// parse path parameter again (as it already contains escaped path information
 	pURL, err := url.Parse(p)
 	if err != nil {
@@ -217,5 +217,5 @@ func buildRequest(p, m string, endpoint url.URL, payload []byte) (*http.Request,
 		endpoint.RawPath = joinPath(endpoint.EscapedPath(), pURL.EscapedPath())
 	}
 	endpoint.Path = joinPath(endpoint.Path, pURL.Path)
-	return http.NewRequest(m, endpoint.String(), bytes.NewBuffer(payload))
+	return http.NewRequestWithContext(ctx, m, endpoint.String(), bytes.NewBuffer(payload))
 }
diff --git a/pkg/client/client_test.go b/pkg/client/client_test.go
@@ -1,6 +1,7 @@
 package client
 
 import (
+	"context"
 	"net/http"
 	"net/url"
 	"testing"
@@ -86,7 +87,7 @@ func TestBuildURL(t *testing.T) {
 			url, err := url.Parse(tt.url)
 			require.NoError(t, err)
 
-			req, err := buildRequest(tt.path, tt.method, *url, []byte{})
+			req, err := buildRequest(context.Background(), tt.path, tt.method, *url, []byte{})
 			require.NoError(t, err)
 			require.Equal(t, tt.resultURL, req.URL.String())
 		})

diff --git a/pkg/client/rules.go b/pkg/client/rules.go
@@ -23,7 +23,7 @@ func (r *CortexClient) CreateRuleGroup(ctx context.Context, namespace string, rg
 	escapedNamespace := url.PathEscape(namespace)
 	path := r.apiPath + "/" + escapedNamespace
 
-	res, err := r.doRequest(path, "POST", payload)
+	res, err := r.doRequest(ctx, path, "POST", payload)
 	if err != nil {
 		return err
 	}
@@ -39,7 +39,7 @@ func (r *CortexClient) DeleteRuleGroup(ctx context.Context, namespace, groupName
 	escapedGroupName := url.PathEscape(groupName)
 	path := r.apiPath + "/" + escapedNamespace + "/" + escapedGroupName
 
-	res, err := r.doRequest(path, "DELETE", nil)
+	res, err := r.doRequest(ctx, path, "DELETE", nil)
 	if err != nil {
 		return err
 	}
@@ -56,7 +56,7 @@ func (r *CortexClient) GetRuleGroup(ctx context.Context, namespace, groupName st
 	path := r.apiPath + "/" + escapedNamespace + "/" + escapedGroupName
 
 	fmt.Println(path)
-	res, err := r.doRequest(path, "GET", nil)
+	res, err := r.doRequest(ctx, path, "GET", nil)
 	if err != nil {
 		return nil, err
 	}
@@ -88,7 +88,7 @@ func (r *CortexClient) ListRules(ctx context.Context, namespace string) (map[str
 		path = path + "/" + namespace
 	}
 
-	res, err := r.doRequest(path, "GET", nil)
+	res, err := r.doRequest(ctx, path, "GET", nil)
 	if err != nil {
 		return nil, err
 	}