Skip to content

Commit ea9610f

Browse files
jackcjackc
authored
Merge pull request #2084 from EpicStep/multiple-tracing
Implement 'MultiTracer'
2 parents 9530aea + 7af618e commit ea9610f

File tree

3 files changed

+268
-1
lines changed

3 files changed

+268
-1
lines changed

‎doc.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ notification is received or the context is canceled.
175175
176176
Tracing and Logging
177177
178-
pgx supports tracing by setting ConnConfig.Tracer.
178+
pgx supports tracing by setting ConnConfig.Tracer. To combine several tracers you can use the multitracer.Tracer.
179179
180180
In addition, the tracelog package provides the TraceLog type which lets a traditional logger act as a Tracer.
181181

‎multitracer/tracer.go

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
// Package multitracer provides a Tracer that can combine several tracers into one.
2+
package multitracer
3+
4+
import (
5+
"context"
6+
7+
"github.com/jackc/pgx/v5"
8+
"github.com/jackc/pgx/v5/pgxpool"
9+
)
10+
11+
// Tracer can combine several tracers into one.
12+
// You can use New to automatically split tracers by interface.
13+
type Tracer struct {
14+
QueryTracers []pgx.QueryTracer
15+
BatchTracers []pgx.BatchTracer
16+
CopyFromTracers []pgx.CopyFromTracer
17+
PrepareTracers []pgx.PrepareTracer
18+
ConnectTracers []pgx.ConnectTracer
19+
PoolAcquireTracers []pgxpool.AcquireTracer
20+
PoolReleaseTracers []pgxpool.ReleaseTracer
21+
}
22+
23+
// New returns new Tracer from tracers with automatically split tracers by interface.
24+
func New(tracers ...pgx.QueryTracer) *Tracer {
25+
var t Tracer
26+
27+
for _, tracer := range tracers {
28+
t.QueryTracers = append(t.QueryTracers, tracer)
29+
30+
if batchTracer, ok := tracer.(pgx.BatchTracer); ok {
31+
t.BatchTracers = append(t.BatchTracers, batchTracer)
32+
}
33+
34+
if copyFromTracer, ok := tracer.(pgx.CopyFromTracer); ok {
35+
t.CopyFromTracers = append(t.CopyFromTracers, copyFromTracer)
36+
}
37+
38+
if prepareTracer, ok := tracer.(pgx.PrepareTracer); ok {
39+
t.PrepareTracers = append(t.PrepareTracers, prepareTracer)
40+
}
41+
42+
if connectTracer, ok := tracer.(pgx.ConnectTracer); ok {
43+
t.ConnectTracers = append(t.ConnectTracers, connectTracer)
44+
}
45+
46+
if poolAcquireTracer, ok := tracer.(pgxpool.AcquireTracer); ok {
47+
t.PoolAcquireTracers = append(t.PoolAcquireTracers, poolAcquireTracer)
48+
}
49+
50+
if poolReleaseTracer, ok := tracer.(pgxpool.ReleaseTracer); ok {
51+
t.PoolReleaseTracers = append(t.PoolReleaseTracers, poolReleaseTracer)
52+
}
53+
}
54+
55+
return &t
56+
}
57+
58+
func (t *Tracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
59+
for _, tracer := range t.QueryTracers {
60+
ctx = tracer.TraceQueryStart(ctx, conn, data)
61+
}
62+
63+
return ctx
64+
}
65+
66+
func (t *Tracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
67+
for _, tracer := range t.QueryTracers {
68+
tracer.TraceQueryEnd(ctx, conn, data)
69+
}
70+
}
71+
72+
func (t *Tracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context {
73+
for _, tracer := range t.BatchTracers {
74+
ctx = tracer.TraceBatchStart(ctx, conn, data)
75+
}
76+
77+
return ctx
78+
}
79+
80+
func (t *Tracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
81+
for _, tracer := range t.BatchTracers {
82+
tracer.TraceBatchQuery(ctx, conn, data)
83+
}
84+
}
85+
86+
func (t *Tracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
87+
for _, tracer := range t.BatchTracers {
88+
tracer.TraceBatchEnd(ctx, conn, data)
89+
}
90+
}
91+
92+
func (t *Tracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
93+
for _, tracer := range t.CopyFromTracers {
94+
ctx = tracer.TraceCopyFromStart(ctx, conn, data)
95+
}
96+
97+
return ctx
98+
}
99+
100+
func (t *Tracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
101+
for _, tracer := range t.CopyFromTracers {
102+
tracer.TraceCopyFromEnd(ctx, conn, data)
103+
}
104+
}
105+
106+
func (t *Tracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
107+
for _, tracer := range t.PrepareTracers {
108+
ctx = tracer.TracePrepareStart(ctx, conn, data)
109+
}
110+
111+
return ctx
112+
}
113+
114+
func (t *Tracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
115+
for _, tracer := range t.PrepareTracers {
116+
tracer.TracePrepareEnd(ctx, conn, data)
117+
}
118+
}
119+
120+
func (t *Tracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context {
121+
for _, tracer := range t.ConnectTracers {
122+
ctx = tracer.TraceConnectStart(ctx, data)
123+
}
124+
125+
return ctx
126+
}
127+
128+
func (t *Tracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
129+
for _, tracer := range t.ConnectTracers {
130+
tracer.TraceConnectEnd(ctx, data)
131+
}
132+
}
133+
134+
func (t *Tracer) TraceAcquireStart(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context {
135+
for _, tracer := range t.PoolAcquireTracers {
136+
ctx = tracer.TraceAcquireStart(ctx, pool, data)
137+
}
138+
139+
return ctx
140+
}
141+
142+
func (t *Tracer) TraceAcquireEnd(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData) {
143+
for _, tracer := range t.PoolAcquireTracers {
144+
tracer.TraceAcquireEnd(ctx, pool, data)
145+
}
146+
}
147+
148+
func (t *Tracer) TraceRelease(pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
149+
for _, tracer := range t.PoolReleaseTracers {
150+
tracer.TraceRelease(pool, data)
151+
}
152+
}

‎multitracer/tracer_test.go

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
package multitracer_test
2+
3+
import (
4+
"context"
5+
"testing"
6+
7+
"github.com/jackc/pgx/v5"
8+
"github.com/jackc/pgx/v5/multitracer"
9+
"github.com/jackc/pgx/v5/pgxpool"
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
type testFullTracer struct{}
14+
15+
func (tt *testFullTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
16+
return ctx
17+
}
18+
19+
func (tt *testFullTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
20+
}
21+
22+
func (tt *testFullTracer) TraceBatchStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchStartData) context.Context {
23+
return ctx
24+
}
25+
26+
func (tt *testFullTracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchQueryData) {
27+
}
28+
29+
func (tt *testFullTracer) TraceBatchEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceBatchEndData) {
30+
}
31+
32+
func (tt *testFullTracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
33+
return ctx
34+
}
35+
36+
func (tt *testFullTracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
37+
}
38+
39+
func (tt *testFullTracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareStartData) context.Context {
40+
return ctx
41+
}
42+
43+
func (tt *testFullTracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) {
44+
}
45+
46+
func (tt *testFullTracer) TraceConnectStart(ctx context.Context, data pgx.TraceConnectStartData) context.Context {
47+
return ctx
48+
}
49+
50+
func (tt *testFullTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEndData) {
51+
}
52+
53+
func (tt *testFullTracer) TraceAcquireStart(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireStartData) context.Context {
54+
return ctx
55+
}
56+
57+
func (tt *testFullTracer) TraceAcquireEnd(ctx context.Context, pool *pgxpool.Pool, data pgxpool.TraceAcquireEndData) {
58+
}
59+
60+
func (tt *testFullTracer) TraceRelease(pool *pgxpool.Pool, data pgxpool.TraceReleaseData) {
61+
}
62+
63+
type testCopyTracer struct{}
64+
65+
func (tt *testCopyTracer) TraceQueryStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryStartData) context.Context {
66+
return ctx
67+
}
68+
69+
func (tt *testCopyTracer) TraceQueryEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceQueryEndData) {
70+
}
71+
72+
func (tt *testCopyTracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context {
73+
return ctx
74+
}
75+
76+
func (tt *testCopyTracer) TraceCopyFromEnd(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromEndData) {
77+
}
78+
79+
func TestNew(t *testing.T) {
80+
t.Parallel()
81+
82+
fullTracer := &testFullTracer{}
83+
copyTracer := &testCopyTracer{}
84+
85+
mt := multitracer.New(fullTracer, copyTracer)
86+
require.Equal(
87+
t,
88+
&multitracer.Tracer{
89+
QueryTracers: []pgx.QueryTracer{
90+
fullTracer,
91+
copyTracer,
92+
},
93+
BatchTracers: []pgx.BatchTracer{
94+
fullTracer,
95+
},
96+
CopyFromTracers: []pgx.CopyFromTracer{
97+
fullTracer,
98+
copyTracer,
99+
},
100+
PrepareTracers: []pgx.PrepareTracer{
101+
fullTracer,
102+
},
103+
ConnectTracers: []pgx.ConnectTracer{
104+
fullTracer,
105+
},
106+
PoolAcquireTracers: []pgxpool.AcquireTracer{
107+
fullTracer,
108+
},
109+
PoolReleaseTracers: []pgxpool.ReleaseTracer{
110+
fullTracer,
111+
},
112+
},
113+
mt,
114+
)
115+
}

0 commit comments

Comments
 (0)