Merge pull request #2084 from EpicStep/multiple-tracing

Implement 'MultiTracer'

doc.go

@@ -175,7 +175,7 @@ notification is received or the context is canceled.
 
 Tracing and Logging
 
-pgx supports tracing by setting ConnConfig.Tracer.
+pgx supports tracing by setting ConnConfig.Tracer. To combine several tracers you can use the multitracer.Tracer.
 
 In addition, the tracelog package provides the TraceLog type which lets a traditional logger act as a Tracer.
 

multitracer/tracer.go

@@ -0,0 +1,152 @@
+// Package multitracer provides a Tracer that can combine several tracers into one.
+package multitracer
+
+import (
+	"context"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+)
+
+// Tracer can combine several tracers into one.
+// You can use New to automatically split tracers by interface.
+type Tracer struct {
+	QueryTracers       []pgx.QueryTracer
+	BatchTracers       []pgx.BatchTracer
+	CopyFromTracers    []pgx.CopyFromTracer
+	PrepareTracers     []pgx.PrepareTracer
+	ConnectTracers     []pgx.ConnectTracer
+	PoolAcquireTracers []pgxpool.AcquireTracer
+	PoolReleaseTracers []pgxpool.ReleaseTracer
+}
+
+// New returns new Tracer from tracers with automatically split tracers by interface.
+func New(tracers ...pgx.QueryTracer) *Tracer {
+	var t Tracer
+
+	for _, tracer := range tracers {
+		t.QueryTracers = append(t.QueryTracers, tracer)
+
+		if batchTracer, ok := tracer.(pgx.BatchTracer); ok {
+			t.BatchTracers = append(t.BatchTracers, batchTracer)
+		}
+
+		if copyFromTracer, ok := tracer.(pgx.CopyFromTracer); ok {
+			t.CopyFromTracers = append(t.CopyFromTracers, copyFromTracer)
+		}
+
+		if prepareTracer, ok := tracer.(pgx.PrepareTracer); ok {
+			t.PrepareTracers = append(t.PrepareTracers, prepareTracer)
+		}
+
+		if connectTracer, ok := tracer.(pgx.ConnectTracer); ok {
+			t.ConnectTracers = append(t.ConnectTracers, connectTracer)
+		}
+
+		if poolAcquireTracer, ok := tracer.(pgxpool.AcquireTracer); ok {
+			t.PoolAcquireTracers = append(t.PoolAcquireTracers, poolAcquireTracer)
+		}
+
+		if poolReleaseTracer, ok := tracer.(pgxpool.ReleaseTracer); ok {
+			t.PoolReleaseTracers = append(t.PoolReleaseTracers, poolReleaseTracer)
+		}
+	}
+
+	return &t
+}
+
+func (t *Tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	for _, tracer := range t.QueryTracers {
+		ctx = tracer.TraceQueryStart(ctx, conn, data)
+	}
+
+	return ctx
+}
+
+func (t *Tracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+	for _, tracer := range t.QueryTracers {
+		tracer.TraceQueryEnd(ctx, conn, data)
+	}
+}
+
+func (t *Tracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context {
+	for _, tracer := range t.BatchTracers {
+		ctx = tracer.TraceBatchStart(ctx, conn, data)
+	}
+
+	return ctx
+}
+
+func (t *Tracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
+	for _, tracer := range t.BatchTracers {
+		tracer.TraceBatchQuery(ctx, conn, data)
+	}
+}
+
+func (t *Tracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
+	for _, tracer := range t.BatchTracers {
+		tracer.TraceBatchEnd(ctx, conn, data)
+	}
+}
+
+func (t *Tracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
+	for _, tracer := range t.CopyFromTracers {
+		ctx = tracer.TraceCopyFromStart(ctx, conn, data)
+	}
+
+	return ctx
+}
+
+func (t *Tracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
+	for _, tracer := range t.CopyFromTracers {
+		tracer.TraceCopyFromEnd(ctx, conn, data)
+	}
+}
+
+func (t *Tracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
+	for _, tracer := range t.PrepareTracers {
+		ctx = tracer.TracePrepareStart(ctx, conn, data)
+	}
+
+	return ctx
+}
+
+func (t *Tracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
+	for _, tracer := range t.PrepareTracers {
+		tracer.TracePrepareEnd(ctx, conn, data)
+	}
+}
+
+func (t *Tracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context {
+	for _, tracer := range t.ConnectTracers {
+		ctx = tracer.TraceConnectStart(ctx, data)
+	}
+
+	return ctx
+}
+
+func (t *Tracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
+	for _, tracer := range t.ConnectTracers {
+		tracer.TraceConnectEnd(ctx, data)
+	}
+}
+
+func (t *Tracer) TraceAcquireStart(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context {
+	for _, tracer := range t.PoolAcquireTracers {
+		ctx = tracer.TraceAcquireStart(ctx, pool, data)
+	}
+
+	return ctx
+}
+
+func (t *Tracer) TraceAcquireEnd(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData) {
+	for _, tracer := range t.PoolAcquireTracers {
+		tracer.TraceAcquireEnd(ctx, pool, data)
+	}
+}
+
+func (t *Tracer) TraceRelease(pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
+	for _, tracer := range t.PoolReleaseTracers {
+		tracer.TraceRelease(pool, data)
+	}
+}

multitracer/tracer_test.go

@@ -0,0 +1,115 @@
+package multitracer_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/multitracer"
+	"github.com/jackc/pgx/v5/pgxpool"
+	"github.com/stretchr/testify/require"
+)
+
+type testFullTracer struct{}
+
+func (tt *testFullTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	return ctx
+}
+
+func (tt *testFullTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+}
+
+func (tt *testFullTracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context {
+	return ctx
+}
+
+func (tt *testFullTracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
+}
+
+func (tt *testFullTracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
+}
+
+func (tt *testFullTracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
+	return ctx
+}
+
+func (tt *testFullTracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
+}
+
+func (tt *testFullTracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
+	return ctx
+}
+
+func (tt *testFullTracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
+}
+
+func (tt *testFullTracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context {
+	return ctx
+}
+
+func (tt *testFullTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
+}
+
+func (tt *testFullTracer) TraceAcquireStart(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context {
+	return ctx
+}
+
+func (tt *testFullTracer) TraceAcquireEnd(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData) {
+}
+
+func (tt *testFullTracer) TraceRelease(pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
+}
+
+type testCopyTracer struct{}
+
+func (tt *testCopyTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
+	return ctx
+}
+
+func (tt *testCopyTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
+}
+
+func (tt *testCopyTracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
+	return ctx
+}
+
+func (tt *testCopyTracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
+}
+
+func TestNew(t *testing.T) {
+	t.Parallel()
+
+	fullTracer := &testFullTracer{}
+	copyTracer := &testCopyTracer{}
+
+	mt := multitracer.New(fullTracer, copyTracer)
+	require.Equal(
+		t,
+		&multitracer.Tracer{
+			QueryTracers: []pgx.QueryTracer{
+				fullTracer,
+				copyTracer,
+			},
+			BatchTracers: []pgx.BatchTracer{
+				fullTracer,
+			},
+			CopyFromTracers: []pgx.CopyFromTracer{
+				fullTracer,
+				copyTracer,
+			},
+			PrepareTracers: []pgx.PrepareTracer{
+				fullTracer,
+			},
+			ConnectTracers: []pgx.ConnectTracer{
+				fullTracer,
+			},
+			PoolAcquireTracers: []pgxpool.AcquireTracer{
+				fullTracer,
+			},
+			PoolReleaseTracers: []pgxpool.ReleaseTracer{
+				fullTracer,
+			},
+		},
+		mt,
+	)
+}