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

@@ -18,32 +18,47 @@ import (
 
 // Wrap augments the given DB with tracing.
 func Wrap(db *pg.DB) {
-	log.Debug("contrib/go-pg/pg.v10: Wrapping Database")
-	db.AddQueryHook(&queryHook{})
-}
-
-type queryHook struct{}
+	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)
+	}
 
-// BeforeQuery implements pg.QueryHook.
-func (h *queryHook) BeforeQuery(ctx context.Context, qe *pg.QueryEvent) (context.Context, error) {
-	query, err := qe.UnformattedQuery()
-	if err != nil {
-		query = []byte("unknown")
+	type queryHook struct{}
+	type queryHook struct{
+		cfg *clientConfig
 	}
 
-	opts := []ddtrace.StartSpanOption{
-		tracer.SpanType(ext.SpanTypeSQL),
-		tracer.ResourceName(string(query)),
+	// BeforeQuery implements pg.QueryHook.
+	func (h *queryHook) BeforeQuery(ctx context.Context, qe *pg.QueryEvent) (context.Context, error) {
+		query, err := qe.UnformattedQuery()
+		if err != nil {
+			query = []byte("unknown")
+		}
+
+		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
 	}
-	_, ctx = tracer.StartSpanFromContext(ctx, "go-pg", opts...)
-	return ctx, qe.Err
-}
 
-// AfterQuery implements pg.QueryHook
-func (h *queryHook) AfterQuery(ctx context.Context, qe *pg.QueryEvent) error {
-	if span, ok := tracer.SpanFromContext(ctx); ok {
+	// AfterQuery implements pg.QueryHook
+	func (h *queryHook) AfterQuery(ctx context.Context, qe *pg.QueryEvent) error {
+		if span, ok := tracer.SpanFromContext(ctx); ok {
 		span.Finish(tracer.WithError(qe.Err))
 	}
 
-	return qe.Err
-}
+		return qe.Err
+	}
+

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

@@ -44,6 +44,7 @@ func TestSelect(t *testing.T) {
 	})
 
 	Wrap(conn)
+	Wrap(conn, WithServiceName("go-pg"))
 
 	parentSpan, ctx := tracer.StartSpanFromContext(context.Background(), "http.request",
 		tracer.ServiceName("fake-http-server"),
@@ -65,4 +66,6 @@ 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))
 }
+