contrib/jackc/pgx.v5: acquire release (#2722)

Co-authored-by: Romain Marcadier <[email protected]>
Co-authored-by: Rodrigo Argüello <[email protected]>

Author: 3 people

.github/workflows/appsec.yml

@@ -258,6 +258,7 @@ jobs:
 
   test-app-smoke-tests:
     name: Smoke Tests
+    if: github.event_name != 'pull_request' || github.event.pull_request.head.repo.owner == 'DataDog'
     uses: DataDog/appsec-go-test-app/.github/workflows/smoke-tests.yml@main
     with:
       dd-trace-go-version: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }}

contrib/jackc/pgx.v5/option.go

@@ -14,6 +14,7 @@ type config struct {
 	traceCopyFrom bool
 	tracePrepare  bool
 	traceConnect  bool
+	traceAcquire  bool
 }
 
 func defaultConfig() *config {
@@ -24,6 +25,7 @@ func defaultConfig() *config {
 		traceCopyFrom: true,
 		tracePrepare:  true,
 		traceConnect:  true,
+		traceAcquire:  true,
 	}
 }
 
@@ -57,6 +59,13 @@ func WithTraceCopyFrom(enabled bool) Option {
 	}
 }
 
+// WithTraceAcquire enables tracing pgxpool connection acquire calls.
+func WithTraceAcquire(enabled bool) Option {
+	return func(c *config) {
+		c.traceAcquire = enabled
+	}
+}
+
 // WithTracePrepare enables tracing prepared statements.
 //
 //	conn, err := pgx.Connect(ctx, "postgres://user:[email protected]:5432/dbname", pgx.WithTraceConnect())

contrib/jackc/pgx.v5/pgx_tracer.go

@@ -7,11 +7,13 @@ 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"
+	"github.com/jackc/pgx/v5/pgxpool"
 )
 
 type operationType string
@@ -30,6 +32,7 @@ const (
 	operationTypePrepare                = "Prepare"
 	operationTypeBatch                  = "Batch"
 	operationTypeCopyFrom               = "Copy From"
+	operationTypeAcquire                = "Acquire"
 )
 
 type tracedBatchQuery struct {
@@ -47,11 +50,12 @@ type pgxTracer struct {
 }
 
 var (
-	_ pgx.QueryTracer    = (*pgxTracer)(nil)
-	_ pgx.BatchTracer    = (*pgxTracer)(nil)
-	_ pgx.ConnectTracer  = (*pgxTracer)(nil)
-	_ pgx.PrepareTracer  = (*pgxTracer)(nil)
-	_ pgx.CopyFromTracer = (*pgxTracer)(nil)
+	_ pgx.QueryTracer       = (*pgxTracer)(nil)
+	_ pgx.BatchTracer       = (*pgxTracer)(nil)
+	_ pgx.ConnectTracer     = (*pgxTracer)(nil)
+	_ pgx.PrepareTracer     = (*pgxTracer)(nil)
+	_ pgx.CopyFromTracer    = (*pgxTracer)(nil)
+	_ pgxpool.AcquireTracer = (*pgxTracer)(nil)
 )
 
 func newPgxTracer(opts ...Option) *pgxTracer {
@@ -175,6 +179,23 @@ func (t *pgxTracer) TraceConnectEnd(ctx context.Context, data pgx.TraceConnectEn
 	finishSpan(ctx, data.Err)
 }
 
+func (t *pgxTracer) TraceAcquireStart(ctx context.Context, pool *pgxpool.Pool, _ pgxpool.TraceAcquireStartData) context.Context {
+	if !t.cfg.traceAcquire {
+		return ctx
+	}
+	opts := t.spanOptions(pool.Config().ConnConfig, operationTypeAcquire, "")
+	_, ctx = tracer.StartSpanFromContext(ctx, "pgx.pool.acquire", opts...)
+	return ctx
+}
+
+func (t *pgxTracer) TraceAcquireEnd(ctx context.Context, _ *pgxpool.Pool, data pgxpool.TraceAcquireEndData) {
+	if !t.cfg.traceAcquire {
+		return
+	}
+
+	finishSpan(ctx, data.Err)
+}
+
 func (t *pgxTracer) spanOptions(connConfig *pgx.ConnConfig, op operationType, sqlStatement string, extraOpts ...ddtrace.StartSpanOption) []ddtrace.StartSpanOption {
 	opts := []ddtrace.StartSpanOption{
 		tracer.ServiceName(t.cfg.serviceName),

contrib/jackc/pgx.v5/pgx_tracer_test.go

@@ -231,13 +231,37 @@ func TestCopyFrom(t *testing.T) {
 	assert.Equal(t, ps.SpanID(), s.ParentID())
 }
 
+func TestAcquire(t *testing.T) {
+	mt := mocktracer.Start()
+	defer mt.Stop()
+
+	opts := append(tracingAllDisabled(), WithTraceAcquire(true))
+	runAllOperations(t, opts...)
+
+	spans := mt.FinishedSpans()
+	require.Len(t, spans, 5)
+
+	ps := spans[4]
+	assert.Equal(t, "parent", ps.OperationName())
+	assert.Equal(t, "parent", ps.Tag(ext.ResourceName))
+
+	s := spans[0]
+	assertCommonTags(t, s)
+	assert.Equal(t, "pgx.pool.acquire", s.OperationName())
+	assert.Equal(t, "Acquire", s.Tag(ext.ResourceName))
+	assert.Equal(t, "Acquire", s.Tag("db.operation"))
+	assert.Equal(t, nil, s.Tag(ext.DBStatement))
+	assert.Equal(t, ps.SpanID(), s.ParentID())
+}
+
 func tracingAllDisabled() []Option {
 	return []Option{
 		WithTraceConnect(false),
 		WithTraceQuery(false),
 		WithTracePrepare(false),
 		WithTraceBatch(false),
 		WithTraceCopyFrom(false),
+		WithTraceAcquire(false),
 	}
 }
 
@@ -246,9 +270,9 @@ func runAllOperations(t *testing.T, opts ...Option) {
 	defer parent.Finish()
 
 	// Connect
-	conn, err := Connect(ctx, postgresDSN, opts...)
+	conn, err := NewPool(ctx, postgresDSN, opts...)
 	require.NoError(t, err)
-	defer conn.Close(ctx)
+	defer conn.Close()
 
 	// Query
 	var x int

contrib/jackc/pgx.v5/pgxpool_test.go

@@ -36,5 +36,5 @@ func TestPool(t *testing.T) {
 	assert.Equal(t, 2, x)
 
 	assert.Len(t, mt.OpenSpans(), 0)
-	assert.Len(t, mt.FinishedSpans(), 5)
+	assert.Len(t, mt.FinishedSpans(), 7)
 }

go.mod

@@ -63,7 +63,7 @@ require (
 	github.com/hashicorp/go-secure-stdlib/parseutil v0.1.7
 	github.com/hashicorp/vault/api v1.9.2
 	github.com/hashicorp/vault/sdk v0.9.2
-	github.com/jackc/pgx/v5 v5.5.4
+	github.com/jackc/pgx/v5 v5.6.0
 	github.com/jinzhu/gorm v1.9.16
 	github.com/jmoiron/sqlx v1.3.5
 	github.com/julienschmidt/httprouter v1.3.0

go.sum

@@ -1517,8 +1517,8 @@ github.com/jackc/pgservicefile v0.0.0-20200714003250-2b9c44734f2b/go.mod h1:vsD4
 github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
 github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
 github.com/jackc/pgx/v5 v5.2.0/go.mod h1:Ptn7zmohNsWEsdxRawMzk3gaKma2obW+NWTnKa0S4nk=
-github.com/jackc/pgx/v5 v5.5.4 h1:Xp2aQS8uXButQdnCMWNmvx6UysWQQC+u1EoizjguY+8=
-github.com/jackc/pgx/v5 v5.5.4/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
+github.com/jackc/pgx/v5 v5.6.0 h1:SWJzexBzPL5jb0GEsrPMLIsi/3jOo7RHlzTjcAeDrPY=
+github.com/jackc/pgx/v5 v5.6.0/go.mod h1:DNZ/vlrUnhWCoFGxHAG8U2ljioxukquj7utPDgtQdTw=
 github.com/jackc/puddle/v2 v2.1.2/go.mod h1:2lpufsF5mRHO6SuZkm0fNYxM6SWHfvyFj62KwNzgels=
 github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
 github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=