From 7af618e42329f7287eee221951fc4b39ec30d26e Mon Sep 17 00:00:00 2001 From: Stepan Rabotkin Date: Sat, 13 Jul 2024 17:04:04 +0300 Subject: [PATCH] feat: add pool tracing to 'MultiTracer' & move it to another package --- doc.go | 2 +- multitracer/tracer.go | 152 +++++++++++++++++++++++++++++++++++++ multitracer/tracer_test.go | 115 ++++++++++++++++++++++++++++ tracer.go | 114 +--------------------------- tracer_test.go | 48 ------------ 5 files changed, 269 insertions(+), 162 deletions(-) create mode 100644 multitracer/tracer.go create mode 100644 multitracer/tracer_test.go diff --git a/doc.go b/doc.go index 8ca07edaf..0e91d64e8 100644 --- a/doc.go +++ b/doc.go @@ -175,7 +175,7 @@ notification is received or the context is canceled. Tracing and Logging -pgx supports tracing by setting ConnConfig.Tracer. To combine several tracers you can use the MultiTracer. +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. diff --git a/multitracer/tracer.go b/multitracer/tracer.go new file mode 100644 index 000000000..acff17398 --- /dev/null +++ b/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) + } +} diff --git a/multitracer/tracer_test.go b/multitracer/tracer_test.go new file mode 100644 index 000000000..aa5ccd080 --- /dev/null +++ b/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, + ) +} diff --git a/tracer.go b/tracer.go index 13c1f8ee2..58ca99f7e 100644 --- a/tracer.go +++ b/tracer.go @@ -2,6 +2,7 @@ package pgx import ( "context" + "github.com/jackc/pgx/v5/pgconn" ) @@ -104,116 +105,3 @@ type TraceConnectEndData struct { Conn *Conn Err error } - -// MultiTracer can combine several tracers into one. -// You can use NewMultiTracer to automatically split tracers by interface. -type MultiTracer struct { - QueryTracers []QueryTracer - BatchTracers []BatchTracer - CopyFromTracers []CopyFromTracer - PrepareTracers []PrepareTracer - ConnectTracers []ConnectTracer -} - -// NewMultiTracer returns new MultiTracer from tracers with automatically split tracers by interface. -func NewMultiTracer(tracers ...QueryTracer) *MultiTracer { - var t MultiTracer - - for _, tracer := range tracers { - t.QueryTracers = append(t.QueryTracers, tracer) - - if batchTracer, ok := tracer.(BatchTracer); ok { - t.BatchTracers = append(t.BatchTracers, batchTracer) - } - - if copyFromTracer, ok := tracer.(CopyFromTracer); ok { - t.CopyFromTracers = append(t.CopyFromTracers, copyFromTracer) - } - - if prepareTracer, ok := tracer.(PrepareTracer); ok { - t.PrepareTracers = append(t.PrepareTracers, prepareTracer) - } - - if connectTracer, ok := tracer.(ConnectTracer); ok { - t.ConnectTracers = append(t.ConnectTracers, connectTracer) - } - } - - return &t -} - -func (t *MultiTracer) TraceQueryStart(ctx context.Context, conn *Conn, data TraceQueryStartData) context.Context { - for _, tracer := range t.QueryTracers { - ctx = tracer.TraceQueryStart(ctx, conn, data) - } - - return ctx -} - -func (t *MultiTracer) TraceQueryEnd(ctx context.Context, conn *Conn, data TraceQueryEndData) { - for _, tracer := range t.QueryTracers { - tracer.TraceQueryEnd(ctx, conn, data) - } -} - -func (t *MultiTracer) TraceBatchStart(ctx context.Context, conn *Conn, data TraceBatchStartData) context.Context { - for _, tracer := range t.BatchTracers { - ctx = tracer.TraceBatchStart(ctx, conn, data) - } - - return ctx -} - -func (t *MultiTracer) TraceBatchQuery(ctx context.Context, conn *Conn, data TraceBatchQueryData) { - for _, tracer := range t.BatchTracers { - tracer.TraceBatchQuery(ctx, conn, data) - } -} - -func (t *MultiTracer) TraceBatchEnd(ctx context.Context, conn *Conn, data TraceBatchEndData) { - for _, tracer := range t.BatchTracers { - tracer.TraceBatchEnd(ctx, conn, data) - } -} - -func (t *MultiTracer) TraceCopyFromStart(ctx context.Context, conn *Conn, data TraceCopyFromStartData) context.Context { - for _, tracer := range t.CopyFromTracers { - ctx = tracer.TraceCopyFromStart(ctx, conn, data) - } - - return ctx -} - -func (t *MultiTracer) TraceCopyFromEnd(ctx context.Context, conn *Conn, data TraceCopyFromEndData) { - for _, tracer := range t.CopyFromTracers { - tracer.TraceCopyFromEnd(ctx, conn, data) - } -} - -func (t *MultiTracer) TracePrepareStart(ctx context.Context, conn *Conn, data TracePrepareStartData) context.Context { - for _, tracer := range t.PrepareTracers { - ctx = tracer.TracePrepareStart(ctx, conn, data) - } - - return ctx -} - -func (t *MultiTracer) TracePrepareEnd(ctx context.Context, conn *Conn, data TracePrepareEndData) { - for _, tracer := range t.PrepareTracers { - tracer.TracePrepareEnd(ctx, conn, data) - } -} - -func (t *MultiTracer) TraceConnectStart(ctx context.Context, data TraceConnectStartData) context.Context { - for _, tracer := range t.ConnectTracers { - ctx = tracer.TraceConnectStart(ctx, data) - } - - return ctx -} - -func (t *MultiTracer) TraceConnectEnd(ctx context.Context, data TraceConnectEndData) { - for _, tracer := range t.ConnectTracers { - tracer.TraceConnectEnd(ctx, data) - } -} diff --git a/tracer_test.go b/tracer_test.go index 82e4d386f..a0fea71e6 100644 --- a/tracer_test.go +++ b/tracer_test.go @@ -566,51 +566,3 @@ func TestTraceConnect(t *testing.T) { require.True(t, traceConnectStartCalled) require.True(t, traceConnectEndCalled) } - -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 TestNewMultiTracer(t *testing.T) { - t.Parallel() - - fullTracer := &testTracer{} - copyTracer := &testCopyTracer{} - - mt := pgx.NewMultiTracer(fullTracer, copyTracer) - require.Equal( - t, - &pgx.MultiTracer{ - QueryTracers: []pgx.QueryTracer{ - fullTracer, - copyTracer, - }, - BatchTracers: []pgx.BatchTracer{ - fullTracer, - }, - CopyFromTracers: []pgx.CopyFromTracer{ - fullTracer, - copyTracer, - }, - PrepareTracers: []pgx.PrepareTracer{ - fullTracer, - }, - ConnectTracers: []pgx.ConnectTracer{ - fullTracer, - }, - }, - mt, - ) -}