Repository: census-ecosystem/opencensus-go-exporter-stackdriver
Pull request #274 (2020-08-14 07:19:50)

Add option to allow user agent to be explicitly overridden


Patch:
diff --git a/stackdriver.go b/stackdriver.go
index fc96955..1e253f9 100644
--- a/stackdriver.go
+++ b/stackdriver.go
@@ -60,6 +60,7 @@ import (
 	metadataapi "cloud.google.com/go/compute/metadata"
 	traceapi "cloud.google.com/go/trace/apiv2"
 	"contrib.go.opencensus.io/exporter/stackdriver/monitoredresource"
+	opencensus "go.opencensus.io"
 	"go.opencensus.io/resource"
 	"go.opencensus.io/resource/resourcekeys"
 	"go.opencensus.io/stats/view"
@@ -279,12 +280,18 @@ type Options struct {
 	// time-series then it will result into an error for the entire CreateTimeSeries request
 	// which may contain more than one time-series.
 	ResourceByDescriptor func(*metricdata.Descriptor, map[string]string) (map[string]string, monitoredresource.Interface)
+
+	// Override the user agent value supplied to Monitoring APIs and included as an
+	// attribute in trace data.
+	UserAgent string
 }
 
 const defaultTimeout = 5 * time.Second
 
 var defaultDomain = path.Join("custom.googleapis.com", "opencensus")
 
+var defaultUserAgent = fmt.Sprintf("opencensus-go %s; stackdriver-exporter %s", opencensus.Version(), version)
+
 // Exporter is a stats and trace exporter that uploads data to Stackdriver.
 //
 // You can create a single Exporter and register it as both a trace exporter
@@ -362,6 +369,9 @@ func NewExporter(o Options) (*Exporter, error) {
 	if o.MetricPrefix != "" && !strings.HasSuffix(o.MetricPrefix, "/") {
 		o.MetricPrefix = o.MetricPrefix + "/"
 	}
+	if o.UserAgent == "" {
+		o.UserAgent = defaultUserAgent
+	}
 
 	se, err := newStatsExporter(o)
 	if err != nil {
diff --git a/stats.go b/stats.go
index 7dde158..67c3c7d 100644
--- a/stats.go
+++ b/stats.go
@@ -26,7 +26,6 @@ import (
 	"sync"
 	"time"
 
-	opencensus "go.opencensus.io"
 	"go.opencensus.io/stats"
 	"go.opencensus.io/stats/view"
 	"go.opencensus.io/tag"
@@ -52,11 +51,9 @@ const (
 	opencensusTaskKey         = "opencensus_task"
 	opencensusTaskDescription = "Opencensus task identifier"
 	defaultDisplayNamePrefix  = "OpenCensus"
-	version                   = "0.10.0"
+	version                   = "0.12.3"
 )
 
-var userAgent = fmt.Sprintf("opencensus-go %s; stackdriver-exporter %s", opencensus.Version(), version)
-
 // statsExporter exports stats to the Stackdriver Monitoring.
 type statsExporter struct {
 	o Options
@@ -89,7 +86,7 @@ func newStatsExporter(o Options) (*statsExporter, error) {
 		return nil, errBlankProjectID
 	}
 
-	opts := append(o.MonitoringClientOptions, option.WithUserAgent(userAgent))
+	opts := append(o.MonitoringClientOptions, option.WithUserAgent(o.UserAgent))
 	ctx := o.Context
 	if ctx == nil {
 		ctx = context.Background()
diff --git a/trace.go b/trace.go
index ae9fa4e..a106613 100644
--- a/trace.go
+++ b/trace.go
@@ -98,7 +98,7 @@ func newTraceExporterWithClient(o Options, c *tracingclient.Client) *traceExport
 
 // ExportSpan exports a SpanData to Stackdriver Trace.
 func (e *traceExporter) ExportSpan(s *trace.SpanData) {
-	protoSpan := protoFromSpanData(s, e.projectID, e.o.Resource)
+	protoSpan := protoFromSpanData(s, e.projectID, e.o.Resource, e.o.UserAgent)
 	protoSize := proto.Size(protoSpan)
 	err := e.bundler.Add(protoSpan, protoSize)
 	switch err {
@@ -137,7 +137,7 @@ func (e *traceExporter) pushTraceSpans(ctx context.Context, node *commonpb.Node,
 	}
 
 	for _, span := range spans {
-		protoSpans = append(protoSpans, protoFromSpanData(span, e.projectID, res))
+		protoSpans = append(protoSpans, protoFromSpanData(span, e.projectID, res, e.o.UserAgent))
 	}
 
 	req := tracepb.BatchWriteSpansRequest{
diff --git a/trace_proto.go b/trace_proto.go
index 422a980..b44957f 100644
--- a/trace_proto.go
+++ b/trace_proto.go
@@ -44,7 +44,7 @@ const (
 )
 
 // proto returns a protocol buffer representation of a SpanData.
-func protoFromSpanData(s *trace.SpanData, projectID string, mr *monitoredrespb.MonitoredResource) *tracepb.Span {
+func protoFromSpanData(s *trace.SpanData, projectID string, mr *monitoredrespb.MonitoredResource, userAgent string) *tracepb.Span {
 	if s == nil {
 		return nil
 	}
@@ -109,6 +109,10 @@ func protoFromSpanData(s *trace.SpanData, projectID string, mr *monitoredrespb.M
 	// Only set the agent label if it is not already set. That enables the
 	// OpenCensus agent/collector to set the agent label based on the library that
 	// sent the span to the agent.
+	//
+	// We now provide a config option to set the userAgent explicitly, which is
+	// used both here and in request headers when sending metric data, but have
+	// retained this non-override functionality for backwards compatibility.
 	if _, hasAgent := sp.Attributes.AttributeMap[agentLabel]; !hasAgent {
 		sp.Attributes.AttributeMap[agentLabel] = &tracepb.AttributeValue{
 			Value: &tracepb.AttributeValue_StringValue{


Test patch:
diff --git a/stackdriver_test.go b/stackdriver_test.go
index 91773dd..3d6fdf9 100644
--- a/stackdriver_test.go
+++ b/stackdriver_test.go
@@ -136,3 +136,14 @@ func TestGRPC(t *testing.T) {
 
 	client.Single(context.Background(), &testpb.FooRequest{SleepNanos: int64(42 * time.Millisecond)})
 }
+
+func TestUserAgent(t *testing.T) {
+	e, err := NewExporter(Options{UserAgent: "OpenCensus Service"})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if want, got := "OpenCensus Service", e.statsExporter.o.UserAgent; want != got {
+		t.Fatalf("UserAgent = %q; want %q", got, want)
+	}
+}
diff --git a/trace_proto_test.go b/trace_proto_test.go
index f09d3cf..85d5343 100644
--- a/trace_proto_test.go
+++ b/trace_proto_test.go
@@ -109,7 +109,7 @@ func generateSpan() {
 }
 
 func createExpectedSpans() spans {
-	ua := trunc(userAgent, len(userAgent))
+	ua := trunc(defaultUserAgent, len(defaultUserAgent))
 	expectedSpans := spans{
 		&tracepb.Span{
 			DisplayName:             trunc("span0", 128),
@@ -284,7 +284,7 @@ func TestExportTrace(t *testing.T) {
 
 	var spbs spans
 	for _, s := range te.spans {
-		spbs = append(spbs, protoFromSpanData(s, "testproject", nil))
+		spbs = append(spbs, protoFromSpanData(s, "testproject", nil, defaultUserAgent))
 	}
 	sort.Sort(spbs)
 
@@ -396,7 +396,7 @@ func TestExportTraceWithMonitoredResource(t *testing.T) {
 	mr := createGCEInstanceMonitoredResource()
 
 	for _, s := range te.spans {
-		gceSpbs = append(gceSpbs, protoFromSpanData(s, "testproject", mr))
+		gceSpbs = append(gceSpbs, protoFromSpanData(s, "testproject", mr, defaultUserAgent))
 	}
 
 	for _, span := range gceSpbs {
@@ -410,7 +410,7 @@ func TestExportTraceWithMonitoredResource(t *testing.T) {
 	mr = createGKEContainerMonitoredResource()
 
 	for _, s := range te.spans {
-		gkeSpbs = append(gkeSpbs, protoFromSpanData(s, "testproject", mr))
+		gkeSpbs = append(gkeSpbs, protoFromSpanData(s, "testproject", mr, defaultUserAgent))
 	}
 
 	for _, span := range gkeSpbs {
@@ -427,7 +427,7 @@ func TestExportTraceWithMonitoredResource(t *testing.T) {
 	var awsEc2Spbs spans
 	mr = createAWSEC2MonitoredResource()
 	for _, s := range te.spans {
-		awsEc2Spbs = append(awsEc2Spbs, protoFromSpanData(s, "testproject", mr))
+		awsEc2Spbs = append(awsEc2Spbs, protoFromSpanData(s, "testproject", mr, defaultUserAgent))
 	}
 
 	for _, span := range awsEc2Spbs {
@@ -500,7 +500,7 @@ func BenchmarkProto(b *testing.B) {
 	}
 	var x int
 	for i := 0; i < b.N; i++ {
-		s := protoFromSpanData(sd, `testproject`, nil)
+		s := protoFromSpanData(sd, `testproject`, nil, defaultUserAgent)
 		x += len(s.Name)
 	}
 	if x == 0 {
diff --git a/trace_test.go b/trace_test.go
index 9a792b0..494c9e4 100644
--- a/trace_test.go
+++ b/trace_test.go
@@ -156,7 +156,7 @@ func TestTraceSpansBufferMaxBytes(t *testing.T) {
 		exported++
 	}
 	for i := 0; i < 10; i++ {
-		e.ExportSpan(makeSampleSpanData())
+		e.ExportSpan(makeSampleSpanData(""))
 	}
 	close(waitCh)
 	e.Flush()
@@ -165,13 +165,47 @@ func TestTraceSpansBufferMaxBytes(t *testing.T) {
 	}
 }
 
-func makeSampleSpanData() *trace.SpanData {
+func TestTraceSpansUserAgent(t *testing.T) {
+	e := newTraceExporterWithClient(Options{
+		UserAgent: "OpenCensus Service",
+		Context:   context.Background(),
+		Timeout:   10 * time.Millisecond,
+	}, nil)
+
+	var got string
+	// set user-agent attribute based on provided option
+	e.uploadFn = func(spans []*tracepb.Span) {
+		got = spans[0].Attributes.AttributeMap[agentLabel].GetStringValue().Value
+	}
+	e.ExportSpan(makeSampleSpanData(""))
+	e.Flush()
+	if want := "OpenCensus Service"; want != got {
+		t.Fatalf("UserAgent Attribute = %q; want %q", got, want)
+	}
+
+	// if user-agent is already set, do not override
+	e.uploadFn = func(spans []*tracepb.Span) {
+		got = spans[0].Attributes.AttributeMap[agentLabel].GetStringValue().Value
+	}
+	e.ExportSpan(makeSampleSpanData("My Test Application"))
+	e.Flush()
+	if want := "My Test Application"; want != got {
+		t.Fatalf("UserAgent Attribute = %q; want %q", got, want)
+	}
+}
+
+func makeSampleSpanData(userAgent string) *trace.SpanData {
 	sd := &trace.SpanData{
 		Annotations:   make([]trace.Annotation, 32),
 		Links:         make([]trace.Link, 32),
 		MessageEvents: make([]trace.MessageEvent, 128),
 		Attributes:    make(map[string]interface{}),
 	}
+
+	if userAgent != "" {
+		sd.Attributes[agentLabel] = userAgent
+	}
+
 	for i := 0; i < 32; i++ {
 		sd.Attributes[fmt.Sprintf("attribute-%d", i)] = ""
 	}