remote-write: add ability to set User-Agent

Signed-off-by: alexgreenbank <alex.greenbank@grafana.com>

storage/remote/client.go

@@ -63,8 +63,15 @@ const (
 )
 
 var (
-	// UserAgent represents Prometheus version to use for user agent header.
+	// internalUserAgent should not be modified as it is to allow tracking of the
-	UserAgent = fmt.Sprintf("Prometheus/%s", version.Version)
+	// specific Prometheus version in use if UserAgent below is over-ridden.
+	// This value is used to set the X-Prometheus-User-Agent header.
+	internalUserAgent = fmt.Sprintf("Prometheus/%s", version.Version)
+
+	// UserAgent represents Prometheus version to use for the User-Agent header.
+	// This is now modifiable.
+	// It defaults to the internalUserAgent value defined above.
+	UserAgent = internalUserAgent
 
 	remoteWriteContentTypeHeaders = map[config.RemoteWriteProtoMsg]string{
 		config.RemoteWriteProtoMsgV1: appProtoContentType, // Also application/x-protobuf;proto=prometheus.WriteRequest but simplified for compatibility with 1.x spec.
@@ -151,6 +158,11 @@ type ReadClient interface {
 	Read(ctx context.Context, query *prompb.Query, sortSeries bool) (storage.SeriesSet, error)
 }
 
+// SetUserAgent allows the User-Agent header value to be over-ridden.
+func SetUserAgent(userAgent string) {
+	UserAgent = userAgent
+}
+
 // NewReadClient creates a new client for remote read.
 func NewReadClient(name string, conf *ClientConfig) (ReadClient, error) {
 	httpClient, err := config_util.NewClientFromConfig(conf.HTTPClientConfig, "remote_storage_read_client")
@@ -262,6 +274,7 @@ func (c *Client) Store(ctx context.Context, req []byte, attempt int) (WriteRespo
 	httpReq.Header.Add("Content-Encoding", string(c.writeCompression))
 	httpReq.Header.Set("Content-Type", remoteWriteContentTypeHeaders[c.writeProtoMsg])
 	httpReq.Header.Set("User-Agent", UserAgent)
+	httpReq.Header.Set("X-Prometheus-User-Agent", internalUserAgent)
 	if c.writeProtoMsg == config.RemoteWriteProtoMsgV1 {
 		// Compatibility mode for 1.0.
 		httpReq.Header.Set(RemoteWriteVersionHeader, RemoteWriteVersion1HeaderValue)
@@ -363,6 +376,7 @@ func (c *Client) Read(ctx context.Context, query *prompb.Query, sortSeries bool)
 	httpReq.Header.Add("Accept-Encoding", "snappy")
 	httpReq.Header.Set("Content-Type", "application/x-protobuf")
 	httpReq.Header.Set("User-Agent", UserAgent)
+	httpReq.Header.Set("X-Prometheus-User-Agent", internalUserAgent)
 	httpReq.Header.Set("X-Prometheus-Remote-Read-Version", "0.1.0")
 
 	ctx, cancel := context.WithTimeout(ctx, c.timeout)

storage/remote/client_test.go

@@ -90,6 +90,146 @@ func TestStoreHTTPErrorHandling(t *testing.T) {
 	}
 }
 
+func TestReadClientUserAgent(t *testing.T) {
+	tests := []struct {
+		name      string
+		userAgent string
+	}{
+		{
+			name:      "default-no-override",
+			userAgent: "",
+		},
+		{
+			name:      "overridden",
+			userAgent: "ArgleBargle/1.2.3",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var called bool
+			server := httptest.NewServer(
+				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					called = true
+					receivedHeaders := r.Header
+
+					// Check the X-Prometheus-User-Agent header.
+					require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("X-Prometheus-User-Agent"),
+						"expected X-Prometheus-User-Agent header to be default value of %q", internalUserAgent)
+
+					if test.userAgent == "" {
+						// Expect original header value.
+						require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("User-Agent"),
+							"expected User-Agent header to be default value of %q", internalUserAgent)
+					} else {
+						// Expect over-ridden header value.
+						require.Equal(t, []string{test.userAgent}, receivedHeaders.Values("User-Agent"),
+							"expected User-Agent header to be over-ridden value of %q", test.userAgent)
+					}
+					w.Header().Set("Content-Type", "text/plain")
+				}),
+			)
+			defer server.Close()
+
+			u, err := url.Parse(server.URL)
+			require.NoError(t, err)
+
+			conf := &ClientConfig{
+				URL:              &config_util.URL{URL: u},
+				Timeout:          model.Duration(5 * time.Second),
+				ChunkedReadLimit: config.DefaultChunkedReadLimit,
+			}
+
+			// Set the User-Agent.
+			if test.userAgent == "" {
+				// Ensure it is set to the default.
+				SetUserAgent(internalUserAgent)
+			} else {
+				SetUserAgent(test.userAgent)
+			}
+
+			c, err := NewReadClient("test", conf)
+			require.NoError(t, err)
+
+			query := &prompb.Query{}
+
+			_, err = c.Read(context.Background(), query, false)
+
+			require.ErrorContains(t, err, "unsupported content type")
+
+			require.True(t, called, "The remote server wasn't called")
+		})
+	}
+}
+
+func TestWriteClientUserAgent(t *testing.T) {
+	tests := []struct {
+		name      string
+		userAgent string
+	}{
+		{
+			name:      "default-no-override",
+			userAgent: "",
+		},
+		{
+			name:      "overridden",
+			userAgent: "ArgleBargle/1.2.3",
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var called bool
+			server := httptest.NewServer(
+				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					called = true
+					receivedHeaders := r.Header
+
+					// Check the X-Prometheus-User-Agent header.
+					require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("X-Prometheus-User-Agent"),
+						"expected X-Prometheus-User-Agent header to be default value of %q", internalUserAgent)
+
+					if test.userAgent == "" {
+						// Expect original header value.
+						require.Equal(t, []string{internalUserAgent}, receivedHeaders.Values("User-Agent"),
+							"expected User-Agent header to be default value of %q", internalUserAgent)
+					} else {
+						// Expect over-ridden header value.
+						require.Equal(t, []string{test.userAgent}, receivedHeaders.Values("User-Agent"),
+							"expected User-Agent header to be over-ridden value of %q", test.userAgent)
+					}
+					// w.Header().Set("Content-Type", "text/plain")
+				}),
+			)
+			defer server.Close()
+
+			serverURL, err := url.Parse(server.URL)
+			require.NoError(t, err)
+
+			conf := &ClientConfig{
+				URL:     &config_util.URL{URL: serverURL},
+				Timeout: model.Duration(time.Second),
+			}
+
+			// Set the User-Agent.
+			if test.userAgent == "" {
+				// Ensure it is set to the default.
+				SetUserAgent(internalUserAgent)
+			} else {
+				SetUserAgent(test.userAgent)
+			}
+
+			hash, err := toHash(conf)
+			require.NoError(t, err)
+			c, err := NewWriteClient(hash, conf)
+			require.NoError(t, err)
+
+			_, err = c.Store(context.Background(), []byte{}, 0)
+			require.NoError(t, err)
+
+			require.True(t, called, "The remote server wasn't called")
+		})
+	}
+}
+
 func TestClientRetryAfter(t *testing.T) {
 	setupServer := func(statusCode int) *httptest.Server {
 		return httptest.NewServer(