feat: pgx & pgxpool support

cmd/instrument.go

@@ -16,6 +16,7 @@ import (
 	"github.com/newrelic/go-easy-instrumentation/integrations/nrgochi"
 	"github.com/newrelic/go-easy-instrumentation/integrations/nrgrpc"
 	"github.com/newrelic/go-easy-instrumentation/integrations/nrnethttp"
+	"github.com/newrelic/go-easy-instrumentation/integrations/nrpgx5"
 	"github.com/newrelic/go-easy-instrumentation/integrations/nrslog"
 	"github.com/newrelic/go-easy-instrumentation/internal/comment"
 	"github.com/newrelic/go-easy-instrumentation/parser"
@@ -54,6 +55,7 @@ func registerIntegrations(manager *parser.InstrumentationManager) {
 		nrgin.InstrumentGinFunction,
 		nrgrpc.InstrumentGrpcServerMethod,
 		nrslog.InstrumentSlogHandler,
+		nrpgx5.InstrumentPgxHandler,
 	)
 
 	// Stateful tracing functions (ORDER PRESERVED)

integrations/nrpgx5/codegen.go

@@ -0,0 +1,73 @@
+package nrpgx5
+
+import (
+	"go/token"
+
+	"github.com/dave/dst"
+)
+
+const (
+	PgxImportPath     = "github.com/jackc/pgx/v5"
+	PgxPoolImportPath = "github.com/jackc/pgx/v5/pgxpool"
+	Nrpgx5ImportPath  = "github.com/newrelic/go-agent/v3/integrations/nrpgx5"
+
+	// configVar is the name used for the synthesized *Config local. It is consistent
+	// across pgx and pgxpool so that the same name appears in every replacement.
+	configVar = "config"
+)
+
+// CreateParseConfig creates `config, err := <pkg>.ParseConfig(connStr)` where pkg is
+// the package identified by importPath (PgxImportPath or PgxPoolImportPath).
+func CreateParseConfig(connStrExpr dst.Expr, importPath string) *dst.AssignStmt {
+	return &dst.AssignStmt{
+		Lhs: []dst.Expr{dst.NewIdent(configVar), dst.NewIdent("err")},
+		Tok: token.DEFINE,
+		Rhs: []dst.Expr{
+			&dst.CallExpr{
+				Fun:  &dst.Ident{Name: "ParseConfig", Path: importPath},
+				Args: []dst.Expr{dst.Clone(connStrExpr).(dst.Expr)},
+			},
+		},
+	}
+}
+
+// CreateTracerAssignment creates `<receiver>.Tracer = nrpgx5.NewTracer()`.
+// receiver selects the field that holds the tracer hook; pass dst.NewIdent("config")
+// for pgx, or config.ConnConfig for pgxpool.
+func CreateTracerAssignment(receiver dst.Expr) *dst.AssignStmt {
+	return &dst.AssignStmt{
+		Lhs: []dst.Expr{
+			&dst.SelectorExpr{
+				X:   receiver,
+				Sel: dst.NewIdent("Tracer"),
+			},
+		},
+		Tok: token.ASSIGN,
+		Rhs: []dst.Expr{newTracerCall()},
+	}
+}
+
+// CreateConnectWithConfig creates `<varName>, err := <fun>(ctx, config)`. fun is the
+// connect-from-config call expression — pgx.ConnectConfig or pgxpool.NewWithConfig.
+func CreateConnectWithConfig(varName string, ctxExpr dst.Expr, fun dst.Expr) *dst.AssignStmt {
+	return &dst.AssignStmt{
+		Lhs: []dst.Expr{dst.NewIdent(varName), dst.NewIdent("err")},
+		Tok: token.DEFINE,
+		Rhs: []dst.Expr{
+			&dst.CallExpr{
+				Fun: fun,
+				Args: []dst.Expr{
+					dst.Clone(ctxExpr).(dst.Expr),
+					dst.NewIdent(configVar),
+				},
+			},
+		},
+	}
+}
+
+// newTracerCall returns a nrpgx5.NewTracer() call expression.
+func newTracerCall() *dst.CallExpr {
+	return &dst.CallExpr{
+		Fun: &dst.Ident{Name: "NewTracer", Path: Nrpgx5ImportPath},
+	}
+}

integrations/nrpgx5/codegen_test.go

@@ -0,0 +1,167 @@
+package nrpgx5
+
+import (
+	"go/token"
+	"testing"
+
+	"github.com/dave/dst"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCreateParseConfig(t *testing.T) {
+	tests := []struct {
+		name       string
+		importPath string
+	}{
+		{name: "pgx", importPath: PgxImportPath},
+		{name: "pgxpool", importPath: PgxPoolImportPath},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			connStr := &dst.BasicLit{Kind: token.STRING, Value: `"postgres://localhost/mydb"`}
+			got := CreateParseConfig(connStr, tt.importPath)
+
+			assert.NotNil(t, got)
+			assert.Equal(t, token.DEFINE, got.Tok)
+			assert.Len(t, got.Lhs, 2)
+
+			configIdent, ok := got.Lhs[0].(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, configVar, configIdent.Name)
+
+			errIdent, ok := got.Lhs[1].(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, "err", errIdent.Name)
+
+			call, ok := got.Rhs[0].(*dst.CallExpr)
+			assert.True(t, ok)
+			assert.Len(t, call.Args, 1)
+
+			funIdent, ok := call.Fun.(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, "ParseConfig", funIdent.Name)
+			assert.Equal(t, tt.importPath, funIdent.Path)
+
+			// Connection string is cloned, not aliased — original must be untouched.
+			arg, ok := call.Args[0].(*dst.BasicLit)
+			assert.True(t, ok)
+			assert.Equal(t, connStr.Value, arg.Value)
+			assert.NotSame(t, connStr, arg)
+		})
+	}
+}
+
+func TestCreateTracerAssignment(t *testing.T) {
+	tests := []struct {
+		name        string
+		receiver    dst.Expr
+		assertOnLhs func(t *testing.T, lhs *dst.SelectorExpr)
+	}{
+		{
+			name:     "direct config receiver (pgx)",
+			receiver: dst.NewIdent(configVar),
+			assertOnLhs: func(t *testing.T, lhs *dst.SelectorExpr) {
+				assert.Equal(t, "Tracer", lhs.Sel.Name)
+				ident, ok := lhs.X.(*dst.Ident)
+				assert.True(t, ok)
+				assert.Equal(t, configVar, ident.Name)
+			},
+		},
+		{
+			name: "nested config.ConnConfig receiver (pgxpool)",
+			receiver: &dst.SelectorExpr{
+				X:   dst.NewIdent(configVar),
+				Sel: dst.NewIdent("ConnConfig"),
+			},
+			assertOnLhs: func(t *testing.T, lhs *dst.SelectorExpr) {
+				assert.Equal(t, "Tracer", lhs.Sel.Name)
+				inner, ok := lhs.X.(*dst.SelectorExpr)
+				assert.True(t, ok)
+				assert.Equal(t, "ConnConfig", inner.Sel.Name)
+				ident, ok := inner.X.(*dst.Ident)
+				assert.True(t, ok)
+				assert.Equal(t, configVar, ident.Name)
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := CreateTracerAssignment(tt.receiver)
+
+			assert.NotNil(t, got)
+			assert.Equal(t, token.ASSIGN, got.Tok)
+			assert.Len(t, got.Lhs, 1)
+
+			lhsSel, ok := got.Lhs[0].(*dst.SelectorExpr)
+			assert.True(t, ok)
+			tt.assertOnLhs(t, lhsSel)
+
+			call, ok := got.Rhs[0].(*dst.CallExpr)
+			assert.True(t, ok)
+
+			funIdent, ok := call.Fun.(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, "NewTracer", funIdent.Name)
+			assert.Equal(t, Nrpgx5ImportPath, funIdent.Path)
+		})
+	}
+}
+
+func TestCreateConnectWithConfig(t *testing.T) {
+	tests := []struct {
+		name       string
+		varName    string
+		fun        dst.Expr
+		wantMethod string
+		wantPath   string
+	}{
+		{
+			name:       "pgx.ConnectConfig",
+			varName:    "conn",
+			fun:        &dst.Ident{Name: "ConnectConfig", Path: PgxImportPath},
+			wantMethod: "ConnectConfig",
+			wantPath:   PgxImportPath,
+		},
+		{
+			name:       "pgxpool.NewWithConfig",
+			varName:    "pool",
+			fun:        &dst.Ident{Name: "NewWithConfig", Path: PgxPoolImportPath},
+			wantMethod: "NewWithConfig",
+			wantPath:   PgxPoolImportPath,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctxExpr := &dst.Ident{Name: "ctx"}
+			got := CreateConnectWithConfig(tt.varName, ctxExpr, tt.fun)
+
+			assert.NotNil(t, got)
+			assert.Equal(t, token.DEFINE, got.Tok)
+			assert.Len(t, got.Lhs, 2)
+
+			lhsIdent, ok := got.Lhs[0].(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, tt.varName, lhsIdent.Name)
+
+			errIdent, ok := got.Lhs[1].(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, "err", errIdent.Name)
+
+			call, ok := got.Rhs[0].(*dst.CallExpr)
+			assert.True(t, ok)
+
+			funIdent, ok := call.Fun.(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, tt.wantMethod, funIdent.Name)
+			assert.Equal(t, tt.wantPath, funIdent.Path)
+
+			assert.Len(t, call.Args, 2)
+			configIdent, ok := call.Args[1].(*dst.Ident)
+			assert.True(t, ok)
+			assert.Equal(t, configVar, configIdent.Name)
+		})
+	}
+}

integrations/nrpgx5/example/pgx/expect.ref

@@ -0,0 +1,35 @@
+--- a/main.go
++++ b/main.go
+@@ -3,15 +3,24 @@
+ import (
+ 	"context"
+ 	"fmt"
++	"time"
+
+ 	"github.com/jackc/pgx/v5"
++	"github.com/newrelic/go-agent/v3/integrations/nrpgx5"
++	"github.com/newrelic/go-agent/v3/newrelic"
+ )
+
+ func main() {
+ 	// Set up a local postgres docker container with:
+ 	// docker run -it -p 5432:5432 -e POSTGRES_PASSWORD=password postgres
++	NewRelicAgent, agentInitError := newrelic.NewApplication(newrelic.ConfigFromEnvironment())
++	if agentInitError != nil {
++		panic(agentInitError)
++	}
+
+-	conn, err := pgx.Connect(context.Background(), "postgres://postgres:password@localhost/postgres")
++	config, err := pgx.ParseConfig("postgres://postgres:password@localhost/postgres")
++	config.Tracer = nrpgx5.NewTracer()
++	conn, err := pgx.ConnectConfig(context.Background(), config)
+ 	if err != nil {
+ 		panic(err)
+ 	}
+@@ -23,4 +24,6 @@
+ 		panic(err)
+ 	}
+ 	fmt.Println(version)
++
++	NewRelicAgent.Shutdown(5 * time.Second)
+ }

integrations/nrpgx5/example/pgx/go.mod

@@ -0,0 +1,11 @@
+module pgx
+
+go 1.25.0
+
+require github.com/jackc/pgx/v5 v5.9.2
+
+require (
+	github.com/jackc/pgpassfile v1.0.0 // indirect
+	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	golang.org/x/text v0.33.0 // indirect
+)

integrations/nrpgx5/example/pgx/go.sum

@@ -0,0 +1,26 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
+github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg=
+github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 h1:iCEnooe7UlwOQYpKFhBabPMi4aNAfoODPEFNiAnClxo=
+github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761/go.mod h1:5TJZWKEWniPve33vlWYSoGYefn3gLQRzjfDlhSJ9ZKM=
+github.com/jackc/pgx/v5 v5.9.2 h1:3ZhOzMWnR4yJ+RW1XImIPsD1aNSz4T4fyP7zlQb56hw=
+github.com/jackc/pgx/v5 v5.9.2/go.mod h1:mal1tBGAFfLHvZzaYh77YS/eC6IX9OWbRV1QIIM0Jn4=
+github.com/jackc/puddle/v2 v2.2.2 h1:PR8nw+E/1w0GLuRFSmiioY6UooMp6KJv0/61nB7icHo=
+github.com/jackc/puddle/v2 v2.2.2/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
+github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
+golang.org/x/sync v0.19.0 h1:vV+1eWNmZ5geRlYjzm2adRgW2/mcpevXNg50YZtPCE4=
+golang.org/x/sync v0.19.0/go.mod h1:9KTHXmSnoGruLpwFjVSX0lNNA75CykiMECbovNTZqGI=
+golang.org/x/text v0.33.0 h1:B3njUFyqtHDUI5jMn1YIr5B0IE2U0qck04r6d4KPAxE=
+golang.org/x/text v0.33.0/go.mod h1:LuMebE6+rBincTi9+xWTY8TztLzKHc/9C1uBCG27+q8=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

integrations/nrpgx5/example/pgx/main.go

@@ -0,0 +1,26 @@
+package main
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/jackc/pgx/v5"
+)
+
+func main() {
+	// Set up a local postgres docker container with:
+	// docker run -it -p 5432:5432 -e POSTGRES_PASSWORD=password postgres
+
+	conn, err := pgx.Connect(context.Background(), "postgres://postgres:password@localhost/postgres")
+	if err != nil {
+		panic(err)
+	}
+	defer conn.Close(context.Background())
+
+	var version string
+	err = conn.QueryRow(context.Background(), "SELECT version()").Scan(&version)
+	if err != nil {
+		panic(err)
+	}
+	fmt.Println(version)
+}

integrations/nrpgx5/example/pgxpool/expect.ref

@@ -0,0 +1,35 @@
+--- a/main.go
++++ b/main.go
+@@ -3,15 +3,24 @@
+ import (
+ 	"context"
+ 	"fmt"
++	"time"
+
+ 	"github.com/jackc/pgx/v5/pgxpool"
++	"github.com/newrelic/go-agent/v3/integrations/nrpgx5"
++	"github.com/newrelic/go-agent/v3/newrelic"
+ )
+
+ func main() {
+ 	// Set up a local postgres docker container with:
+ 	// docker run -it -p 5432:5432 -e POSTGRES_PASSWORD=password postgres
++	NewRelicAgent, agentInitError := newrelic.NewApplication(newrelic.ConfigFromEnvironment())
++	if agentInitError != nil {
++		panic(agentInitError)
++	}
+
+-	pool, err := pgxpool.New(context.Background(), "postgres://postgres:password@localhost/postgres")
++	config, err := pgxpool.ParseConfig("postgres://postgres:password@localhost/postgres")
++	config.ConnConfig.Tracer = nrpgx5.NewTracer()
++	pool, err := pgxpool.NewWithConfig(context.Background(), config)
+ 	if err != nil {
+ 		panic(err)
+ 	}
+@@ -23,4 +24,6 @@
+ 		panic(err)
+ 	}
+ 	fmt.Println(version)
++
++	NewRelicAgent.Shutdown(5 * time.Second)
+ }