From a78763098891d39dc457f62518cc33e1644493ca Mon Sep 17 00:00:00 2001 From: Stepan Rabotkin Date: Thu, 11 Jul 2024 23:41:17 +0300 Subject: [PATCH 1/4] feat: implement 'MultiTracer' --- tracer.go | 110 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) diff --git a/tracer.go b/tracer.go index 58ca99f7..29635c17 100644 --- a/tracer.go +++ b/tracer.go @@ -105,3 +105,113 @@ type TraceConnectEndData struct { Conn *Conn Err error } + +type MultiTracer struct { + queryTracers []QueryTracer + batchTracers []BatchTracer + copyFromTracers []CopyFromTracer + prepareTracers []PrepareTracer + connectTracers []ConnectTracer +} + +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) + } +} From 8e46d2117c149650e40f692d0b706df8be8e8698 Mon Sep 17 00:00:00 2001 From: Stepan Rabotkin Date: Sat, 13 Jul 2024 01:40:46 +0300 Subject: [PATCH 2/4] refac: export 'MultiTracer' fields --- tracer.go | 43 +++++++++++++++++++++---------------------- 1 file changed, 21 insertions(+), 22 deletions(-) diff --git a/tracer.go b/tracer.go index 29635c17..8d3fb12e 100644 --- a/tracer.go +++ b/tracer.go @@ -2,7 +2,6 @@ package pgx import ( "context" - "github.com/jackc/pgx/v5/pgconn" ) @@ -107,33 +106,33 @@ type TraceConnectEndData struct { } type MultiTracer struct { - queryTracers []QueryTracer - batchTracers []BatchTracer - copyFromTracers []CopyFromTracer - prepareTracers []PrepareTracer - connectTracers []ConnectTracer + QueryTracers []QueryTracer + BatchTracers []BatchTracer + CopyFromTracers []CopyFromTracer + PrepareTracers []PrepareTracer + ConnectTracers []ConnectTracer } func NewMultiTracer(tracers ...QueryTracer) *MultiTracer { var t MultiTracer for _, tracer := range tracers { - t.queryTracers = append(t.queryTracers, tracer) + t.QueryTracers = append(t.QueryTracers, tracer) if batchTracer, ok := tracer.(BatchTracer); ok { - t.batchTracers = append(t.batchTracers, batchTracer) + t.BatchTracers = append(t.BatchTracers, batchTracer) } if copyFromTracer, ok := tracer.(CopyFromTracer); ok { - t.copyFromTracers = append(t.copyFromTracers, copyFromTracer) + t.CopyFromTracers = append(t.CopyFromTracers, copyFromTracer) } if prepareTracer, ok := tracer.(PrepareTracer); ok { - t.prepareTracers = append(t.prepareTracers, prepareTracer) + t.PrepareTracers = append(t.PrepareTracers, prepareTracer) } if connectTracer, ok := tracer.(ConnectTracer); ok { - t.connectTracers = append(t.connectTracers, connectTracer) + t.ConnectTracers = append(t.ConnectTracers, connectTracer) } } @@ -141,7 +140,7 @@ func NewMultiTracer(tracers ...QueryTracer) *MultiTracer { } func (t *MultiTracer) TraceQueryStart(ctx context.Context, conn *Conn, data TraceQueryStartData) context.Context { - for _, tracer := range t.queryTracers { + for _, tracer := range t.QueryTracers { ctx = tracer.TraceQueryStart(ctx, conn, data) } @@ -149,13 +148,13 @@ func (t *MultiTracer) TraceQueryStart(ctx context.Context, conn *Conn, data Trac } func (t *MultiTracer) TraceQueryEnd(ctx context.Context, conn *Conn, data TraceQueryEndData) { - for _, tracer := range t.queryTracers { + 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 { + for _, tracer := range t.BatchTracers { ctx = tracer.TraceBatchStart(ctx, conn, data) } @@ -163,19 +162,19 @@ func (t *MultiTracer) TraceBatchStart(ctx context.Context, conn *Conn, data Trac } func (t *MultiTracer) TraceBatchQuery(ctx context.Context, conn *Conn, data TraceBatchQueryData) { - for _, tracer := range t.batchTracers { + 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 { + 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 { + for _, tracer := range t.CopyFromTracers { ctx = tracer.TraceCopyFromStart(ctx, conn, data) } @@ -183,13 +182,13 @@ func (t *MultiTracer) TraceCopyFromStart(ctx context.Context, conn *Conn, data T } func (t *MultiTracer) TraceCopyFromEnd(ctx context.Context, conn *Conn, data TraceCopyFromEndData) { - for _, tracer := range t.copyFromTracers { + 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 { + for _, tracer := range t.PrepareTracers { ctx = tracer.TracePrepareStart(ctx, conn, data) } @@ -197,13 +196,13 @@ func (t *MultiTracer) TracePrepareStart(ctx context.Context, conn *Conn, data Tr } func (t *MultiTracer) TracePrepareEnd(ctx context.Context, conn *Conn, data TracePrepareEndData) { - for _, tracer := range t.prepareTracers { + 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 { + for _, tracer := range t.ConnectTracers { ctx = tracer.TraceConnectStart(ctx, data) } @@ -211,7 +210,7 @@ func (t *MultiTracer) TraceConnectStart(ctx context.Context, data TraceConnectSt } func (t *MultiTracer) TraceConnectEnd(ctx context.Context, data TraceConnectEndData) { - for _, tracer := range t.connectTracers { + for _, tracer := range t.ConnectTracers { tracer.TraceConnectEnd(ctx, data) } } From 3f270eec7d1dbb0b9c0fa427f76c5c23800936a2 Mon Sep 17 00:00:00 2001 From: Stepan Rabotkin Date: Sat, 13 Jul 2024 02:02:22 +0300 Subject: [PATCH 3/4] feat: add 'MultiTracer' to go doc & cover it --- doc.go | 2 +- tracer.go | 3 +++ tracer_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 1 deletion(-) diff --git a/doc.go b/doc.go index bc0391dd..8ca07eda 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. +pgx supports tracing by setting ConnConfig.Tracer. To combine several tracers you can use the MultiTracer. In addition, the tracelog package provides the TraceLog type which lets a traditional logger act as a Tracer. diff --git a/tracer.go b/tracer.go index 8d3fb12e..13c1f8ee 100644 --- a/tracer.go +++ b/tracer.go @@ -105,6 +105,8 @@ type TraceConnectEndData struct { 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 @@ -113,6 +115,7 @@ type MultiTracer struct { ConnectTracers []ConnectTracer } +// NewMultiTracer returns new MultiTracer from tracers with automatically split tracers by interface. func NewMultiTracer(tracers ...QueryTracer) *MultiTracer { var t MultiTracer diff --git a/tracer_test.go b/tracer_test.go index a0fea71e..82e4d386 100644 --- a/tracer_test.go +++ b/tracer_test.go @@ -566,3 +566,51 @@ 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, + ) +} From 7af618e42329f7287eee221951fc4b39ec30d26e Mon Sep 17 00:00:00 2001 From: Stepan Rabotkin Date: Sat, 13 Jul 2024 17:04:04 +0300 Subject: [PATCH 4/4] 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 8ca07eda..0e91d64e 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 00000000..acff1739 --- /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 00000000..aa5ccd08 --- /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 13c1f8ee..58ca99f7 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 82e4d386..a0fea71e 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, - ) -}