Repository: newrelic/newrelic-client-go
Pull request #985 (2023-02-06 21:56:42)

feat: facilitate additional service names for NewRelic-Requesting-Services request header via env var
This PR enables prepending additional requesting service names to a request header to facilitate tracking additional services that call the client. For example, the CLI we'll be able to know the CLI was called from with our Deployment Marker GitHub Action.


Patch:
diff --git a/internal/http/client.go b/internal/http/client.go
index 51bca82..7825a88 100644
--- a/internal/http/client.go
+++ b/internal/http/client.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"os"
 	"reflect"
 	"regexp"
 	"strings"
@@ -76,6 +77,16 @@ func NewClient(cfg config.Config) Client {
 		cfg.ServiceName = fmt.Sprintf("%s|%s", cfg.ServiceName, defaultServiceName)
 	}
 
+	// If a requesting service sets the NEW_RELIC_SERVICE_NAME env variable,
+	// we prepend this additional custom service name to the existing service name.
+	// The service name is used to track which requesting service is being utilized.
+	// e.g. We can track usage of the New Relic Deployment Marker GitHub Action by
+	// setting this environment variable when executing a command.
+	customServiceName := os.Getenv("NEW_RELIC_SERVICE_NAME")
+	if customServiceName != "" {
+		cfg.ServiceName = fmt.Sprintf("%s|%s", customServiceName, cfg.ServiceName)
+	}
+
 	r := retryablehttp.NewClient()
 	r.HTTPClient = &c
 	r.RetryMax = defaultRetryMax
diff --git a/internal/http/request.go b/internal/http/request.go
index 320fb0e..6e40a97 100644
--- a/internal/http/request.go
+++ b/internal/http/request.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"os"
 
 	"github.com/google/go-querystring/query"
 	retryablehttp "github.com/hashicorp/go-retryablehttp"
@@ -121,6 +122,12 @@ func (r *Request) SetErrorValue(e ErrorResponse) {
 // SetServiceName sets the service name for the request.
 func (r *Request) SetServiceName(serviceName string) {
 	serviceName = fmt.Sprintf("%s|%s", serviceName, defaultServiceName)
+
+	customServiceName := os.Getenv("NEW_RELIC_SERVICE_NAME")
+	if customServiceName != "" {
+		serviceName = fmt.Sprintf("%s|%s", customServiceName, serviceName)
+	}
+
 	r.SetHeader(defaultNewRelicRequestingServiceHeader, serviceName)
 }
 
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 4cfd7c5..3aa3024 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -2,7 +2,9 @@
 package config
 
 import (
+	"fmt"
 	"net/http"
+	"os"
 	"time"
 
 	"github.com/newrelic/newrelic-client-go/v2/internal/version"
@@ -90,6 +92,19 @@ func (c *Config) SetRegion(reg *region.Region) error {
 	return nil
 }
 
+func (c *Config) SetServiceName(serviceName string) error {
+	if c.ServiceName == "" {
+		c.ServiceName = serviceName
+	}
+
+	customServiceName := os.Getenv("NEW_RELIC_SERVICE_NAME")
+	if customServiceName != "" {
+		c.ServiceName = fmt.Sprintf("%s|%s", customServiceName, c.ServiceName)
+	}
+
+	return nil
+}
+
 // GetLogger returns a logger instance based on the config values.
 func (c *Config) GetLogger() logging.Logger {
 	if c.Logger != nil {


Test patch:
diff --git a/internal/http/client_test.go b/internal/http/client_test.go
index 36f6042..651deb6 100644
--- a/internal/http/client_test.go
+++ b/internal/http/client_test.go
@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
+	"os"
 	"testing"
 	"time"
 
@@ -316,6 +317,37 @@ func TestCustomRequestHeaders(t *testing.T) {
 	_, err = c.Do(req)
 
 	assert.Nil(t, err)
+	assert.Equal(t, "custom-requesting-service|newrelic-client-go", req.GetHeader("NewRelic-Requesting-Services"))
+}
+
+func TestCustomRequestHeadersWithEnvironmentVariable(t *testing.T) {
+	// Set environment variable for testing
+	_ = os.Setenv("NEW_RELIC_SERVICE_NAME", "testing")
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		w.WriteHeader(http.StatusOK)
+
+		assert.Equal(t, "custom-user-agent", r.Header.Get("user-agent"))
+		assert.Equal(t, "testing|custom-requesting-service|newrelic-client-go", r.Header.Get("newrelic-requesting-services"))
+	}))
+
+	tc := mock.NewTestConfig(t, ts)
+
+	c := NewClient(tc)
+
+	req, err := c.NewRequest("GET", c.config.Region().RestURL("path"), nil, nil, nil)
+
+	req.SetHeader("user-agent", "custom-user-agent")
+	req.SetServiceName("custom-requesting-service")
+
+	_, err = c.Do(req)
+
+	assert.Nil(t, err)
+	assert.Equal(t, "testing|custom-requesting-service|newrelic-client-go", req.GetHeader("NewRelic-Requesting-Services"))
+
+	// Unset environment variable to prevent polluting the rest of the test run
+	_ = os.Unsetenv("NEW_RELIC_SERVICE_NAME")
 }
 
 func TestAccountIDHeaderWithPersonalAPIKeyCapableV2Authorizer(t *testing.T) {