contrib/go-pg/pg.v10 add options, ServiceName

contrib/go-pg/pg.v10/option.go

@@ -0,0 +1,60 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2016 Datadog, Inc.
+
+package pg
+
+import (
+	"math"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/internal"
+)
+
+type clientConfig struct {
+	serviceName   string
+	analyticsRate float64
+}
+
+// ClientOption represents an option that can be used to create or wrap a client.
+type ClientOption func(*clientConfig)
+
+func defaults(cfg *clientConfig) {
+	cfg.serviceName = "gopg.db"
+	// cfg.analyticsRate = globalconfig.AnalyticsRate()
+	if internal.BoolEnv("DD_TRACE_GOPG_ANALYTICS_ENABLED", false) {
+		cfg.analyticsRate = 1.0
+	} else {
+		cfg.analyticsRate = math.NaN()
+	}
+}
+
+// WithServiceName sets the given service name for the client.
+func WithServiceName(name string) ClientOption {
+	return func(cfg *clientConfig) {
+		cfg.serviceName = name
+	}
+}
+
+// WithAnalytics enables Trace Analytics for all started spans.
+func WithAnalytics(on bool) ClientOption {
+	return func(cfg *clientConfig) {
+		if on {
+			cfg.analyticsRate = 1.0
+		} else {
+			cfg.analyticsRate = math.NaN()
+		}
+	}
+}
+
+// WithAnalyticsRate sets the sampling rate for Trace Analytics events
+// correlated to started spans.
+func WithAnalyticsRate(rate float64) ClientOption {
+	return func(cfg *clientConfig) {
+		if rate >= 0.0 && rate <= 1.0 {
+			cfg.analyticsRate = rate
+		} else {
+			cfg.analyticsRate = math.NaN()
+		}
+	}
+}

contrib/go-pg/pg.v10/pg_go.go

@@ -17,12 +17,22 @@ import (
 )
 
 // Wrap augments the given DB with tracing.
-func Wrap(db *pg.DB) {
+func Wrap(db *pg.DB, opts ...ClientOption) {
+	cfg := new(clientConfig)
+	defaults(cfg)
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	h := &queryHook{
+		cfg: cfg,
+	}
 	log.Debug("contrib/go-pg/pg.v10: Wrapping Database")
-	db.AddQueryHook(&queryHook{})
+	db.AddQueryHook(h)
 }
 
-type queryHook struct{}
+type queryHook struct {
+	cfg *clientConfig
+}
 
 // BeforeQuery implements pg.QueryHook.
 func (h *queryHook) BeforeQuery(ctx context.Context, qe *pg.QueryEvent) (context.Context, error) {
@@ -34,6 +44,7 @@ func (h *queryHook) BeforeQuery(ctx context.Context, qe *pg.QueryEvent) (context
 	opts := []ddtrace.StartSpanOption{
 		tracer.SpanType(ext.SpanTypeSQL),
 		tracer.ResourceName(string(query)),
+		tracer.ServiceName(h.cfg.serviceName),
 	}
 	_, ctx = tracer.StartSpanFromContext(ctx, "go-pg", opts...)
 	return ctx, qe.Err

contrib/go-pg/pg.v10/pg_go_test.go

@@ -43,7 +43,7 @@ func TestSelect(t *testing.T) {
 		Database: "postgres",
 	})
 
-	Wrap(conn)
+	Wrap(conn, WithServiceName("go-pg"))
 
 	parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
 		tracer.ServiceName("fake-http-server"),
@@ -65,4 +65,5 @@ func TestSelect(t *testing.T) {
 	assert.Equal(1, n)
 	assert.Equal("go-pg", spans[0].OperationName())
 	assert.Equal("http.request", spans[1].OperationName())
+	assert.Equal("go-pg", span.Tag(ext.ServiceName))
 }