diff --git a/v3/integrations/nrpgx5/nrpgx5.go b/v3/integrations/nrpgx5/nrpgx5.go index 6331064fb..9d2fdc0c7 100644 --- a/v3/integrations/nrpgx5/nrpgx5.go +++ b/v3/integrations/nrpgx5/nrpgx5.go @@ -16,36 +16,36 @@ // // For example: // -// import ( -// "github.com/jackc/pgx/v5" -// "github.com/newrelic/go-agent/v3/integrations/nrpgx5" -// "github.com/newrelic/go-agent/v3/newrelic" -// ) -// -// func main() { -// cfg, err := pgx.ParseConfig("postgres://postgres:postgres@localhost:5432") // OR pgxpools.ParseConfig(...) -// if err != nil { -// panic(err) -// } -// -// cfg.Tracer = nrpgx5.NewTracer() -// conn, err := pgx.ConnectConfig(context.Background(), cfg) -// if err != nil { -// panic(err) -// } -// } +// import ( +// "github.com/jackc/pgx/v5" +// "github.com/newrelic/go-agent/v3/integrations/nrpgx5" +// "github.com/newrelic/go-agent/v3/newrelic" +// ) +// +// func main() { +// cfg, err := pgx.ParseConfig("postgres://postgres:postgres@localhost:5432") // OR pgxpools.ParseConfig(...) +// if err != nil { +// panic(err) +// } +// +// cfg.Tracer = nrpgx5.NewTracer() +// conn, err := pgx.ConnectConfig(context.Background(), cfg) +// if err != nil { +// panic(err) +// } +// } // // See the programs in the example directory for working examples of each use case. package nrpgx5 import ( "context" - "strconv" - "github.com/jackc/pgx/v5" "github.com/newrelic/go-agent/v3/internal" "github.com/newrelic/go-agent/v3/newrelic" "github.com/newrelic/go-agent/v3/newrelic/sqlparse" + "strconv" + "strings" ) func init() { @@ -74,14 +74,16 @@ type TracerOption func(*Tracer) // NewTracer creates a new value which implements pgx.BatchTracer, pgx.ConnectTracer, pgx.PrepareTracer, and pgx.QueryTracer. // This value will be used to facilitate instrumentation of the database operations performed. // When establishing a connection to the database, the recommended usage is to do something like the following: -// cfg, err := pgx.ParseConfig("...") -// if err != nil { ... } -// cfg.Tracer = nrpgx5.NewTracer() -// conn, err := pgx.ConnectConfig(context.Background(), cfg) +// +// cfg, err := pgx.ParseConfig("...") +// if err != nil { ... } +// cfg.Tracer = nrpgx5.NewTracer() +// conn, err := pgx.ConnectConfig(context.Background(), cfg) // // If you do not wish to have SQL query parameters included in the telemetry data, add the WithQueryParameters // option, like so: -// cfg.Tracer = nrpgx5.NewTracer(nrpgx5.WithQueryParameters(false)) +// +// cfg.Tracer = nrpgx5.NewTracer(nrpgx5.WithQueryParameters(false)) // // (The default is to collect query parameters, but you can explicitly select this by passing true to WithQueryParameters.) // @@ -222,3 +224,26 @@ func (t *Tracer) TracePrepareStart(ctx context.Context, conn *pgx.Conn, data pgx // TracePrepareEnd implements pgx.PrepareTracer. func (t *Tracer) TracePrepareEnd(ctx context.Context, conn *pgx.Conn, data pgx.TracePrepareEndData) { } + +// TraceCopyFromStart is called at the beginning of CopyFrom calls. The +// returned context is used for the rest of the call and will be passed to +// TraceCopyFromEnd. +func (t *Tracer) TraceCopyFromStart(ctx context.Context, conn *pgx.Conn, data pgx.TraceCopyFromStartData) context.Context { + segment := t.BaseSegment + segment.StartTime = newrelic.FromContext(ctx).StartSegmentNow() + + segment.Operation = "copy_from" + segment.Collection = strings.ReplaceAll(data.TableName.Sanitize(), "\"", "") + segment.AddAttribute("db.columnNames", data.ColumnNames) + + return context.WithValue(ctx, querySegmentKey, &segment) +} + +// TraceCopyFromEnd is called at the end of CopyFrom calls. +func (t *Tracer) TraceCopyFromEnd(ctx context.Context, _ *pgx.Conn, data pgx.TraceCopyFromEndData) { + segment, ok := ctx.Value(querySegmentKey).(*newrelic.DatastoreSegment) + if !ok { + return + } + segment.End() +} diff --git a/v3/integrations/nrpgx5/nrpgx5_test.go b/v3/integrations/nrpgx5/nrpgx5_test.go index 90e6be1b3..32f35a500 100644 --- a/v3/integrations/nrpgx5/nrpgx5_test.go +++ b/v3/integrations/nrpgx5/nrpgx5_test.go @@ -292,6 +292,27 @@ func TestTracer_inPool(t *testing.T) { } } +func TestTracer_copyFrom(t *testing.T) { + conn, finish := getTestCon(t) + defer finish() + + t.Run("copy from should send metric with table identifier", func(t *testing.T) { + app := integrationsupport.NewBasicTestApp() + + txn := app.StartTransaction(t.Name()) + + ctx := newrelic.NewContext(context.Background(), txn) + _, _ = conn.CopyFrom(ctx, pgx.Identifier{"mytable"}, []string{"name"}, pgx.CopyFromRows([][]interface{}{{"name a"}, {"name b"}, {"name c"}})) + + txn.End() + + app.ExpectMetricsPresent(t, []internal.WantMetric{ + {Name: "Datastore/operation/Postgres/copy_from"}, + {Name: "Datastore/statement/Postgres/mytable/copy_from"}, + }) + }) +} + func getTestCon(t testing.TB) (*pgx.Conn, func()) { snap := pgsnap.NewSnap(t, os.Getenv("PGSNAP_DB_URL"))