diff --git a/metrics/newrelic/provider.go b/metrics/newrelic/provider.go index 9f952c7..1a6cfa2 100644 --- a/metrics/newrelic/provider.go +++ b/metrics/newrelic/provider.go @@ -1,6 +1,7 @@ package newrelic import ( + "context" "net/http" "net/url" "time" @@ -104,6 +105,13 @@ func (t *Trace) Unwrap() *newrelic.Transaction { return t.txn } +// EnrichContext implements metrics.ContextEnricher to add the New Relic +// transaction to the context using New Relic's context key. This enables +// automatic instrumentation by libraries like nrpgx. +func (t *Trace) EnrichContext(ctx context.Context) context.Context { + return newrelic.NewContext(ctx, t.txn) +} + // Span wraps a New Relic segment type Span struct { seg *newrelic.Segment diff --git a/metrics/tracing.go b/metrics/tracing.go index af5c265..40e6602 100644 --- a/metrics/tracing.go +++ b/metrics/tracing.go @@ -36,6 +36,14 @@ type Span interface { End() } +// ContextEnricher is an optional interface that traces can implement to add +// provider-specific data to the context. This is used by instrumentation +// libraries (e.g., nrpgx) that look for transactions using their own context keys. +type ContextEnricher interface { + // EnrichContext adds provider-specific data to the context + EnrichContext(ctx context.Context) context.Context +} + // TraceMethodCall traces a method call with a given struct/package and method names func TraceMethodCall(ctx context.Context, structOrPackageName, methodName string) *MethodTracer { trace := TraceFromContext(ctx) @@ -106,9 +114,15 @@ type traceContextKey struct{} // TraceKey is the context key for Trace var TraceKey = traceContextKey{} -// NewContext returns a new context with the trace attached +// NewContext returns a new context with the trace attached. +// If the trace implements ContextEnricher, it also enriches the context +// with provider-specific data for automatic instrumentation libraries. func NewContext(ctx context.Context, trace Trace) context.Context { - return context.WithValue(ctx, TraceKey, trace) + ctx = context.WithValue(ctx, TraceKey, trace) + if enricher, ok := trace.(ContextEnricher); ok { + ctx = enricher.EnrichContext(ctx) + } + return ctx } // TraceFromContext retrieves the trace from context, if present