feat instrument pgx CopyFrom

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()
+}

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"))