WIP: contrib/jackc/pgx.v5: add support for pgx

contrib/jackc/pgx.v5/example_test.go

@@ -0,0 +1,64 @@
+// 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 2022 Datadog, Inc.
+
+package pgx_test
+
+import (
+	"context"
+	"log"
+
+	pgxtrace "gopkg.in/DataDog/dd-trace-go.v1/contrib/jackc/pgx.v5"
+
+	"github.com/jackc/pgx/v5"
+)
+
+func ExampleConnect() {
+	ctx := context.TODO()
+
+	// The package exposes the same connect functions as the regular pgx.v5 library
+	// which sets up a tracer and connects as usual.
+	db, err := pgxtrace.Connect(ctx, "postgres://pqgotest:password@localhost/pqgotest?sslmode=disable")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer db.Close(ctx)
+
+	// Any calls made to the database will be traced as expected.
+	rows, err := db.Query(ctx, "SELECT name FROM users WHERE age=$1", 27)
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer rows.Close()
+
+	// This enables you to use PostgresSQL specific functions implemented by pgx.v5.
+	numbers := []int{1, 2, 3}
+
+	copyFromSource := pgx.CopyFromSlice(len(numbers), func(i int) ([]any, error) {
+		return []any{numbers[i]}, nil
+	})
+
+	_, err = db.CopyFrom(ctx, []string{"numbers"}, []string{"number"}, copyFromSource)
+	if err != nil {
+		log.Fatal(err)
+	}
+}
+
+func ExamplePool() {
+	ctx := context.TODO()
+
+	// The pgxpool uses the same tracer and is exposed the same way.
+	pool, err := pgxtrace.NewPool(ctx, "postgres://pqgotest:password@localhost/pqgotest?sslmode=disable")
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer pool.Close()
+
+	var x int
+
+	err = pool.QueryRow(ctx, "SELECT 1").Scan(&x)
+	if err != nil {
+		log.Fatal(err)
+	}
+}

contrib/jackc/pgx.v5/option.go

@@ -0,0 +1,83 @@
+// 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 2022 Datadog, Inc.
+
+package pgx
+
+import "gopkg.in/DataDog/dd-trace-go.v1/internal/namingschema"
+
+type config struct {
+	serviceName   string
+	traceQuery    bool
+	traceBatch    bool
+	traceCopyFrom bool
+	tracePrepare  bool
+	traceConnect  bool
+}
+
+func defaultConfig() *config {
+	return &config{
+		serviceName:   namingschema.NewDefaultServiceName(defaultServiceName).GetName(),
+		traceQuery:    true,
+		traceBatch:    true,
+		traceCopyFrom: true,
+		tracePrepare:  true,
+		traceConnect:  true,
+	}
+}
+
+type Option func(*config)
+
+// WithServiceName sets the service name to use for all spans.
+func WithServiceName(name string) Option {
+	return func(c *config) {
+		c.serviceName = name
+	}
+}
+
+// WithTraceQuery enables tracing query operations.
+func WithTraceQuery(enabled bool) Option {
+	return func(c *config) {
+		c.traceQuery = enabled
+	}
+}
+
+// WithTraceBatch enables tracing batched operations (i.e. pgx.Batch{}).
+func WithTraceBatch(enabled bool) Option {
+	return func(c *config) {
+		c.traceBatch = enabled
+	}
+}
+
+// WithTraceCopyFrom enables tracing pgx.CopyFrom calls.
+func WithTraceCopyFrom(enabled bool) Option {
+	return func(c *config) {
+		c.traceCopyFrom = enabled
+	}
+}
+
+// WithTracePrepare enables tracing prepared statements.
+//
+//	conn, err := pgx.Connect(ctx, "postgres://user:pass@example.com:5432/dbname", pgx.WithTraceConnect())
+//	if err != nil {
+//		// handle err
+//	}
+//	defer conn.Close(ctx)
+//
+//	_, err := conn.Prepare(ctx, "stmt", "select $1::integer")
+//	row, err := conn.QueryRow(ctx, "stmt", 1)
+func WithTracePrepare(enabled bool) Option {
+	return func(c *config) {
+		c.tracePrepare = enabled
+	}
+}
+
+// WithTraceConnect enables tracing calls to Connect and ConnectConfig.
+//
+//	pgx.Connect(ctx, "postgres://user:pass@example.com:5432/dbname", pgx.WithTraceConnect())
+func WithTraceConnect(enabled bool) Option {
+	return func(c *config) {
+		c.traceConnect = enabled
+	}
+}

contrib/jackc/pgx.v5/pgx.go

@@ -0,0 +1,44 @@
+// 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 2022 Datadog, Inc.
+
+package pgx
+
+import (
+	"context"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+	"gopkg.in/DataDog/dd-trace-go.v1/internal/telemetry"
+
+	"github.com/jackc/pgx/v5"
+)
+
+const (
+	componentName      = "jackc/pgx.v5"
+	defaultServiceName = "postgres.db"
+)
+
+func init() {
+	telemetry.LoadIntegration(componentName)
+	tracer.MarkIntegrationImported("github.com/jackc/pgx.v5")
+}
+
+type Batch = pgx.Batch
+
+func Connect(ctx context.Context, connString string, opts ...Option) (*pgx.Conn, error) {
+	connConfig, err := pgx.ParseConfig(connString)
+	if err != nil {
+		return nil, err
+	}
+
+	return ConnectConfig(ctx, connConfig, opts...)
+}
+
+func ConnectConfig(ctx context.Context, connConfig *pgx.ConnConfig, opts ...Option) (*pgx.Conn, error) {
+	// The tracer must be set in the config before calling connect
+	// as pgx takes ownership of the config. QueryTracer traces
+	// may work, but none of the others will, as they're set in
+	// unexported fields in the config in the pgx.connect function.
+	connConfig.Tracer = newPgxTracer(opts...)
+	return pgx.ConnectConfig(ctx, connConfig)
+}

contrib/jackc/pgx.v5/pgx_tracer.go

@@ -0,0 +1,155 @@
+// 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 2022 Datadog, Inc.
+
+package pgx
+
+import (
+	"context"
+
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer"
+
+	"github.com/jackc/pgx/v5"
+)
+
+type pgxTracer struct {
+	cfg *config
+}
+
+var (
+	_ pgx.QueryTracer    = (*pgxTracer)(nil)
+	_ pgx.BatchTracer    = (*pgxTracer)(nil)
+	_ pgx.ConnectTracer  = (*pgxTracer)(nil)
+	_ pgx.PrepareTracer  = (*pgxTracer)(nil)
+	_ pgx.CopyFromTracer = (*pgxTracer)(nil)
+)
+
+func newPgxTracer(opts ...Option) *pgxTracer {
+	cfg := defaultConfig()
+	for _, opt := range opts {
+		opt(cfg)
+	}
+	return &pgxTracer{cfg: cfg}
+}
+
+func (t *pgxTracer) defaultSpanOptions() []ddtrace.StartSpanOption {
+	return []ddtrace.StartSpanOption{
+		tracer.ServiceName(t.cfg.serviceName),
+		tracer.SpanType(ext.SpanTypeSQL),
+		tracer.Tag(ext.DBSystem, ext.DBSystemPostgreSQL),
+		tracer.Tag(ext.Component, componentName),
+		tracer.Tag(ext.SpanKind, ext.SpanKindClient),
+	}
+}
+
+func (t *pgxTracer) spanOptions(opts ...ddtrace.StartSpanOption) []ddtrace.StartSpanOption {
+	return append(t.defaultSpanOptions(), opts...)
+}
+
+func finish(ctx context.Context, err error) {
+	span, ok := tracer.SpanFromContext(ctx)
+	if !ok {
+		return
+	}
+	span.Finish(tracer.WithError(err))
+}
+
+func (t *pgxTracer) TraceQueryStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	if !t.cfg.traceQuery {
+		return ctx
+	}
+	opts := t.spanOptions(
+		tracer.ResourceName(data.SQL),
+	)
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.query", opts...)
+	return ctx
+}
+
+func (t *pgxTracer) TraceQueryEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceQueryEndData) {
+	if !t.cfg.traceQuery {
+		return
+	}
+	finish(ctx, data.Err)
+}
+
+func (t *pgxTracer) TraceBatchStart(ctx context.Context, _ *pgx.Conn, _ pgx.TraceBatchStartData) context.Context {
+	if !t.cfg.traceBatch {
+		return ctx
+	}
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.batch", t.spanOptions()...)
+	return ctx
+}
+
+func (t *pgxTracer) TraceBatchQuery(ctx context.Context, _ *pgx.Conn, data pgx.TraceBatchQueryData) {
+	if !t.cfg.traceBatch {
+		return
+	}
+	opts := t.spanOptions(
+		tracer.ResourceName(data.SQL),
+	)
+	// TODO: this might be wrong
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.batch.query", opts...)
+	finish(ctx, data.Err)
+}
+
+func (t *pgxTracer) TraceBatchEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceBatchEndData) {
+	if !t.cfg.traceBatch {
+		return
+	}
+	finish(ctx, data.Err)
+}
+
+func (t *pgxTracer) TraceCopyFromStart(ctx context.Context, _ *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
+	if !t.cfg.traceCopyFrom {
+		return ctx
+	}
+	opts := t.spanOptions(
+		tracer.Tag("tables", data.TableName),
+		tracer.Tag("columns", data.ColumnNames),
+	)
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.copyfrom", opts...)
+	return ctx
+}
+
+func (t *pgxTracer) TraceCopyFromEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceCopyFromEndData) {
+	if !t.cfg.traceCopyFrom {
+		return
+	}
+	finish(ctx, data.Err)
+}
+
+func (t *pgxTracer) TracePrepareStart(ctx context.Context, _ *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
+	if !t.cfg.tracePrepare {
+		return ctx
+	}
+	opts := t.spanOptions(
+		tracer.ResourceName(data.SQL),
+	)
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.prepare", opts...)
+	return ctx
+}
+
+func (t *pgxTracer) TracePrepareEnd(ctx context.Context, _ *pgx.Conn, data pgx.TracePrepareEndData) {
+	if !t.cfg.tracePrepare {
+		return
+	}
+	finish(ctx, data.Err)
+}
+
+func (t *pgxTracer) TraceConnectStart(ctx context.Context, _ pgx.TraceConnectStartData) context.Context {
+	if !t.cfg.traceConnect {
+		return ctx
+	}
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.connect", t.spanOptions()...)
+	return ctx
+}
+
+func (t *pgxTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
+	if !t.cfg.traceConnect {
+		return
+	}
+	finish(ctx, data.Err)
+}