mudhorn/prometheus
1
0
Fork 0
mirror of https://github.com/prometheus/prometheus.git synced 2024-12-25 21:54:10 -08:00
Code Issues Actions 3 Packages Projects Releases Wiki Activity

remote-write: add ability to set User-Agent

Browse source 
Signed-off-by: alexgreenbank <alex.greenbank@grafana.com>
This commit is contained in:
alexgreenbank 2024-10-23 13:01:30 +01:00
parent 7c7116fea8
commit 62e6d1abf5
2 changed files with 156 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
storage/remote/client.go
View file

@ -63,8 +63,15 @@ const (
)
var (
// UserAgent represents Prometheus version to use for user agent header.
UserAgent = fmt.Sprintf("Prometheus/%s", version.Version)
// internalUserAgent should not be modified as it is to allow tracking of the
// 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)

140
storage/remote/client_test.go
View file

@ -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(