diff --git a/options.go b/options.go index f0411c2..55e4052 100644 --- a/options.go +++ b/options.go @@ -29,6 +29,26 @@ func WithTracerProvider(provider trace.TracerProvider) Option { }) } +// WithBatchQueryTracerProvider specifies a tracer provider to use for creating a tracer +// that is used specifically for spans within a batch's per-query operations. +// +// This is useful for those who may have lengthy batches and wish to either disable, +// or by some means sample the spans created for each individual query within a batch operation. +// +// For instance, one could choose to sample 50% of the per-query spans created within a batch +// by specifying the following tracer provider to WithBatchQueryTracerProvider: +// trace.NewTracerProvider(trace.WithSampler(trace.ParentBased(trace.TraceIDRatioBased(0.5)))) +// +// If WithBatchQueryTracerProvider is not specified, the same tracer provider as specified by WithTracerProvider is used. +// If WithTracerProvider is not specified either, then the same default as WithTracerProvider is used -- the OTel SDK global provider. +func WithBatchQueryTracerProvider(provider trace.TracerProvider) Option { + return optionFunc(func(cfg *tracerConfig) { + if provider != nil { + cfg.batchQueryTracerProvider = provider + } + }) +} + // WithMeterProvider specifies a meter provider to use for creating a meter. // If none is specified, the global provider is used. func WithMeterProvider(provider metric.MeterProvider) Option { diff --git a/tracer.go b/tracer.go index 670d0f0..951cf6f 100644 --- a/tracer.go +++ b/tracer.go @@ -58,10 +58,11 @@ var _ pgxpool.AcquireTracer = (*Tracer)(nil) // Tracer is a wrapper around the pgx tracer interfaces which instrument // queries with both tracing and metrics. type Tracer struct { - tracer trace.Tracer - meter metric.Meter - tracerAttrs []attribute.KeyValue - meterAttrs []attribute.KeyValue + batchQueryTracer trace.Tracer + tracer trace.Tracer + meter metric.Meter + tracerAttrs []attribute.KeyValue + meterAttrs []attribute.KeyValue operationDuration metric.Int64Histogram operationErrors metric.Int64Counter @@ -75,8 +76,9 @@ type Tracer struct { } type tracerConfig struct { - tracerProvider trace.TracerProvider - meterProvider metric.MeterProvider + batchQueryTracerProvider trace.TracerProvider + tracerProvider trace.TracerProvider + meterProvider metric.MeterProvider tracerAttrs []attribute.KeyValue meterAttrs []attribute.KeyValue @@ -92,8 +94,9 @@ type tracerConfig struct { // NewTracer returns a new Tracer. func NewTracer(opts ...Option) *Tracer { cfg := &tracerConfig{ - tracerProvider: otel.GetTracerProvider(), - meterProvider: otel.GetMeterProvider(), + batchQueryTracerProvider: nil, + tracerProvider: otel.GetTracerProvider(), + meterProvider: otel.GetMeterProvider(), tracerAttrs: []attribute.KeyValue{ semconv.DBSystemPostgreSQL, }, @@ -112,7 +115,16 @@ func NewTracer(opts ...Option) *Tracer { opt.apply(cfg) } + // If a batchQueryTracerProvider was not optionally assigned, + // default to using the same TracerProvider as used for all other operations. + // This is guaranteed to be non-nil as this defaults to the global TracerProvider, + // unless otherwise optionally user-assigned as well. + if cfg.batchQueryTracerProvider == nil { + cfg.batchQueryTracerProvider = cfg.tracerProvider + } + tracer := &Tracer{ + batchQueryTracer: cfg.batchQueryTracerProvider.Tracer(tracerName, trace.WithInstrumentationVersion(findOwnImportedVersion())), tracer: cfg.tracerProvider.Tracer(tracerName, trace.WithInstrumentationVersion(findOwnImportedVersion())), meter: cfg.meterProvider.Meter(meterName, metric.WithInstrumentationVersion(findOwnImportedVersion())), tracerAttrs: cfg.tracerAttrs, @@ -399,7 +411,7 @@ func (t *Tracer) TraceBatchQuery(ctx context.Context, conn *pgx.Conn, data pgx.T } } - _, span := t.tracer.Start(ctx, spanName, opts...) + _, span := t.batchQueryTracer.Start(ctx, spanName, opts...) recordSpanError(span, data.Err) span.End()